Fix relabeling in per-FOV labeling, after switching to ROIs (ref #115)

fractal/tasks/image_labeling.py

@@ -11,7 +11,6 @@
 Institute for Biomedical Research and Pelkmans Lab from the University of
 Zurich.
 """
-import itertools
 import json
 import shutil
 import time
@@ -43,13 +42,10 @@ def segment_FOV(
     logfile="LOG_image_labeling",
 ):
 
-    # chunk_location = block_info[None]["chunk-location"]
-    chunk_location = "dummy"
-
     # Write some debugging info
     with open(logfile, "a") as out:
         out.write(
-            f"[{chunk_location}] START Cellpose |"
+            f"[segment_FOV] START Cellpose |"
             f" column: {type(column)}, {column.shape} |"
             f" do_3D: {do_3D}\n"
         )
@@ -73,7 +69,7 @@ def segment_FOV(
     # Write some debugging info
     with open(logfile, "a") as out:
         out.write(
-            f"[{chunk_location}] END   Cellpose |"
+            f"[segment_FOV] END   Cellpose |"
             f" Elapsed: {t1-t0:.4f} seconds |"
             f" mask shape: {mask.shape},"
             f" mask dtype: {mask.dtype} (before recast to {label_dtype}),"
@@ -93,7 +89,7 @@ def image_labeling(
     relabeling=True,
     anisotropy=None,
     diameter_level0=80.0,
-    cellprob_threshold=None,
+    cellprob_threshold=0.0,
     model_type="nuclei",
 ):
 
@@ -174,10 +170,6 @@ def image_labeling(
                     f"pixel_size_y={pixel_size_y}"
                 )
             anisotropy = pixel_size_z / pixel_size_x
-    else:
-        raise NotImplementedError(
-            "TODO: The integration of 2D labeling with ROIs is not ready yet"
-        )
 
     # Check model_type
     if model_type not in ["nuclei", "cyto2", "cyto"]:
@@ -235,6 +227,7 @@ def image_labeling(
     with open(logfile, "w") as out:
         out.write(f"Start image_labeling task for {zarrurl}\n")
         out.write(f"relabeling: {relabeling}\n")
+        out.write(f"labeling_level: {labeling_level}\n")
         out.write(f"model_type: {model_type}\n")
         out.write(f"anisotropy: {anisotropy}\n")
         out.write(f"num_threads: {num_threads}\n")
@@ -331,12 +324,13 @@ def image_labeling(
         # https://stackoverflow.com/a/72018364/19085332
         num_labels_tot = 0
         num_labels_column = 0
-        for inds in itertools.product(
-            *map(range, mask_rechunked.blocks.shape)
-        ):
+
+        for indices in list_indices:
+            s_z, e_z, s_y, e_y, s_x, e_x = indices[:]
+            shape = [e_z - s_z, e_y - s_y, e_x - s_x]
 
             # Select a specific chunk (=column in 3D, =image in 2D)
-            column_mask = mask_rechunked.blocks[inds].compute()
+            column_mask = mask_rechunked[s_z:e_z, s_y:e_y, s_x:e_x].compute()
             num_labels_column = np.max(column_mask)
 
             # Apply re-labeling and update total number of labels
@@ -346,7 +340,7 @@ def image_labeling(
 
             with open(logfile, "a") as out:
                 out.write(
-                    f"Chunk {inds}, "
+                    f"FOV ROI {indices}, "
                     f"num_labels_column={num_labels_column}, "
                     f"num_labels_tot={num_labels_tot}\n"
                 )
@@ -359,15 +353,7 @@ def image_labeling(
                     f"but dtype={label_dtype}"
                 )
             # Re-assign the chunk to the new array
-            start_z = inds[0] * nz
-            end_z = (inds[0] + 1) * nz
-            start_y = inds[1] * img_size_y
-            end_y = (inds[1] + 1) * img_size_y
-            start_x = inds[2] * img_size_x
-            end_x = (inds[2] + 1) * img_size_x
-            newmask_rechunked[start_z:end_z, start_y:end_y, start_x:end_x] = (
-                column_mask + shift
-            )
+            newmask_rechunked[s_z:e_z, s_y:e_y, s_x:e_x] = column_mask + shift
 
         newmask = newmask_rechunked.rechunk(data_zyx.chunks)
 
@@ -438,6 +424,13 @@ def image_labeling(
         type=int,
         help="TBD",
     )
+    parser.add_argument(
+        "-ll",
+        "--labeling_level",
+        default=0,
+        type=int,
+        help="TBD",
+    )
 
     args = parser.parse_args()
     image_labeling(
@@ -446,5 +439,5 @@ def image_labeling(
         chl_list=args.chl_list,
         labeling_channel=args.labeling_channel,
         num_threads=args.num_threads,
-        # FIXME: more arguments
+        labeling_level=args.labeling_level,
     )