From 6b995e73254695e441b64a67053eb99f4144e92c Mon Sep 17 00:00:00 2001
From: Martin Laasmaa <martin@sysbio.ioc.ee>
Date: Tue, 16 Dec 2025 13:24:20 +0200
Subject: [PATCH] upate

---
 src/model/yolo_wrapper.py     |  4 ++--
 src/utils/image_converters.py | 37 +++++++++++++++++++++++++++--------
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/src/model/yolo_wrapper.py b/src/model/yolo_wrapper.py
index d5e31b3..f08ef69 100644
--- a/src/model/yolo_wrapper.py
+++ b/src/model/yolo_wrapper.py
@@ -9,7 +9,7 @@ from typing import Optional, List, Dict, Callable, Any
 import torch
 import tempfile
 import os
-from src.utils.image import Image, convert_grayscale_to_rgb_preserve_range
+from src.utils.image import Image
 from src.utils.logger import get_logger
 
 
@@ -238,7 +238,7 @@ class YOLOWrapper:
                     img_obj = Image(source_path)
                     pil_img = img_obj.pil_image
                     if len(pil_img.getbands()) == 1:
-                        rgb_img = convert_grayscale_to_rgb_preserve_range(pil_img)
+                        rgb_img = img_obj.convert_grayscale_to_rgb_preserve_range()
                     else:
                         rgb_img = pil_img.convert("RGB")
 
diff --git a/src/utils/image_converters.py b/src/utils/image_converters.py
index 7597ec5..13ac80e 100644
--- a/src/utils/image_converters.py
+++ b/src/utils/image_converters.py
@@ -12,18 +12,27 @@ class UT:
     Operetta files along with rois drawn in ImageJ
     """
 
-    def __init__(self, roifile_fn: Path):
+    def __init__(self, roifile_fn: Path, no_labels: bool):
         self.roifile_fn = roifile_fn
         print("is file", self.roifile_fn.is_file())
-        self.rois = ImagejRoi.fromfile(self.roifile_fn)
-        self.stem = self.roifile_fn.stem.split("Roi-")[1]
+        self.rois = None
+        if no_labels:
+            self.rois = ImagejRoi.fromfile(self.roifile_fn)
+            self.stem = self.roifile_fn.stem.split("Roi-")[1]
+        else:
+            self.roifile_fn = roifile_fn / roifile_fn.parts[-1]
+            self.stem = self.roifile_fn.stem
+
+            print(self.roifile_fn)
+
+        print(self.stem)
         self.image, self.image_props = self._load_images()
 
     def _load_images(self):
         """Loading sequence of tif files
         array sequence is CZYX
         """
-        print(self.roifile_fn.parent, self.stem)
+        print("Loading images:", self.roifile_fn.parent, self.stem)
         fns = list(self.roifile_fn.parent.glob(f"{self.stem.lower()}*.tif*"))
         stems = [fn.stem.split(self.stem)[-1] for fn in fns]
         n_ch = len(set([stem.split("-ch")[-1].split("t")[0] for stem in stems]))
@@ -116,6 +125,7 @@ class UT:
             self.image = np.max(self.image[channel], axis=0)
             print(self.image.shape)
 
+        print(path / subfolder / f"{self.stem}.tif")
         with TiffWriter(path / subfolder / f"{self.stem}.tif") as tif:
             tif.write(self.image)
 
@@ -126,14 +136,25 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("-i", "--input", nargs="*", type=Path)
     parser.add_argument("-o", "--output", type=Path)
+    parser.add_argument(
+        "--no-labels",
+        action="store_false",
+        help="Source does not have labels, export only images",
+    )
     args = parser.parse_args()
 
     for path in args.input:
         print("Path:", path)
-        for rfn in Path(path).glob("*.zip"):
-            print("Roi FN:", rfn)
-            ut = UT(rfn)
-            ut.export_rois(args.output, class_index=0)
+        if not args.no_labels:
+            print("No labels")
+            ut = UT(path, args.no_labels)
             ut.export_image(args.output, plane_mode="max projection", channel=0)
 
+        else:
+            for rfn in Path(path).glob("*.zip"):
+                print("Roi FN:", rfn)
+                ut = UT(rfn, args.no_labels)
+                ut.export_rois(args.output, class_index=0)
+                ut.export_image(args.output, plane_mode="max projection", channel=0)
+
         print()