From ec1de42fe4a871606052ae34ac68728c6a81a7fc Mon Sep 17 00:00:00 2001
From: "Christopher L. Crutchfield" <ccrutchf@ucsd.edu>
Date: Tue, 12 Nov 2024 17:04:55 -0800
Subject: [PATCH] feat: output mask as part of execution for fishial

---
 fishsense_lite/commands/label_studio.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/fishsense_lite/commands/label_studio.py b/fishsense_lite/commands/label_studio.py
index 59f1c99..2000160 100644
--- a/fishsense_lite/commands/label_studio.py
+++ b/fishsense_lite/commands/label_studio.py
@@ -83,6 +83,10 @@ def execute_fishial(
 ):
     device = "cuda" if torch.cuda.is_available() else "cpu"
     output_file = get_output_file(input_file, root, output, "jpg")
+    mask_file = (
+        output_file.parent
+        / f"{output_file.name.removesuffix(output_file.suffix)}.mask.png"
+    )
     json_file = output_file.with_suffix(".json")
 
     if output_file.exists() and json_file.exists() and not overwrite:
@@ -100,6 +104,11 @@ def execute_fishial(
     output_file.parent.mkdir(parents=True, exist_ok=True)
     cv2.imwrite(output_file.absolute().as_posix(), image)
 
+    debug_output = (segmentations.astype(float) / segmentations.max() * 255).astype(
+        np.uint8
+    )
+    cv2.imwrite(mask_file.absolute().as_posix(), debug_output)
+
     json_objects = SegmentationLabelStudioJSON(
         f"{prefix}{output_file.relative_to(output.absolute()).as_posix()}",
         segmentations,