Update nellie.py

nellie.py

@@ -1,6 +1,7 @@
 import os
 import glob
 import logging
+import random
 import numpy as np
 import matplotlib.pyplot as plt
 import tifffile as tiff
@@ -15,7 +16,7 @@ from nellie.segmentation.labelling import Label
 from skimage.filters import rank, sobel, gaussian
 from scipy.ndimage import median_filter
 from skimage import exposure
-from skimage.morphology import disk
+from skimage.morphology import disk, binary_dilation
 from skimage.segmentation import find_boundaries
 from skimage.util import img_as_ubyte
 
@@ -31,12 +32,10 @@ def preprocess_method_1(img_raw: np.ndarray) -> np.ndarray:
     img_min, img_max = np.min(img_float), np.max(img_float)
     img_8bit = (((img_float - img_min) / (img_max - img_min) * 255).astype(np.uint8) 
                 if img_max > img_min else np.zeros(img_raw.shape, dtype=np.uint8))
-
     img_inverted = 255 - img_8bit
     img_blur = gaussian(img_inverted, sigma=1, preserve_range=True)
     img_filtered = median_filter(img_blur, size=5).astype(np.uint8)
     p_low, p_high = np.percentile(img_filtered, (0.75, 99.25))
-
     if p_low >= p_high: return img_filtered
     return exposure.rescale_intensity(img_filtered, in_range=(p_low, p_high), out_range=(0, 255)).astype(np.uint8)
 
@@ -54,7 +53,6 @@ def preprocess_method_2(img_input: np.ndarray) -> np.ndarray:
 def preprocess_combined(img_raw: np.ndarray) -> np.ndarray:
     return preprocess_method_2(preprocess_method_1(img_raw))
 
-
 # =====================================================================
 # Vectorized Metric Evaluation (Fixed IoU and Bounds)
 # =====================================================================
@@ -63,25 +61,23 @@ def evaluate_predictions(pred_mask: np.ndarray, gt_mask: np.ndarray, iou_thresho
     pred_mask, gt_mask = pred_mask.astype(np.int32), gt_mask.astype(np.int32)
     pred_labels = np.unique(pred_mask)[np.unique(pred_mask) != 0]
     gt_labels = np.unique(gt_mask)[np.unique(gt_mask) != 0]
-
     num_pred, num_gt = len(pred_labels), len(gt_labels)
-
+    
     pred_bg, gt_bg = pred_mask > 0, gt_mask > 0
     union_pixel = np.logical_or(pred_bg, gt_bg).sum()
     global_pixel_iou = np.logical_and(pred_bg, gt_bg).sum() / union_pixel if union_pixel > 0 else 0.0
-
+    
     if num_gt == 0 and num_pred == 0: return global_pixel_iou, 1.0, 1.0, 1.0, num_pred, num_gt
     if num_gt == 0 or num_pred == 0: return global_pixel_iou, 0.0, 0.0, 0.0, num_pred, num_gt
-
+    
     overlapping = gt_bg & pred_bg
     if not np.any(overlapping): return global_pixel_iou, 0.0, 0.0, 0.0, num_pred, num_gt
-
+    
     gt_id_map = {id_: i for i, id_ in enumerate(gt_labels)}
     pred_id_map = {id_: j for j, id_ in enumerate(pred_labels)}
-
+    
     gt_indices = np.array([gt_id_map[x] for x in gt_mask[overlapping]])
     pred_indices = np.array([pred_id_map[x] for x in pred_mask[overlapping]])
-
     intersection = csr_matrix((np.ones(len(gt_indices), dtype=np.int32), (gt_indices, pred_indices)), shape=(num_gt, num_pred)).toarray()
     
     gt_vec = np.array([np.bincount(gt_mask.ravel())[id_] for id_ in gt_labels])[:, None]
@@ -89,18 +85,16 @@ def evaluate_predictions(pred_mask: np.ndarray, gt_mask: np.ndarray, iou_thresho
     
     union = gt_vec + pred_vec - intersection
     iou_matrix = np.divide(intersection, union, out=np.zeros_like(intersection, dtype=float), where=union != 0)
-
-    # TP is the number of predictions that successfully match a GT with IoU > threshold
+    
     matches = iou_matrix > iou_threshold
-    tp = np.sum(np.any(matches, axis=0)) # Number of matched predictions
+    tp = np.sum(np.any(matches, axis=0)) 
     
     precision = tp / num_pred if num_pred > 0 else 0.0
     recall = tp / num_gt if num_gt > 0 else 0.0
     instance_f1 = (2 * precision * recall) / (precision + recall) if (precision + recall) > 0 else 0.0
-
+    
     return global_pixel_iou, precision, recall, instance_f1, num_pred, num_gt
 
-
 # =====================================================================
 # Dataset Prep & Core Pipeline
 # =====================================================================
@@ -112,7 +106,7 @@ def prepare_datasets(raw_tif_files, base_dir="generated_datasets"):
         "Prep 1->2": os.path.join(base_dir, "dataset_prep1_2")
     }
     for d in dirs.values(): os.makedirs(d, exist_ok=True)
-
+    
     for filepath in tqdm(raw_tif_files, desc="Building Datasets"):
         fname = os.path.basename(filepath)
         if os.path.exists(os.path.join(dirs["Prep 1->2"], fname)):
@@ -123,17 +117,18 @@ def prepare_datasets(raw_tif_files, base_dir="generated_datasets"):
         tiff.imwrite(os.path.join(dirs["Prep 1"], fname), preprocess_method_1(img_raw))
         tiff.imwrite(os.path.join(dirs["Prep 2"], fname), preprocess_method_2(img_raw))
         tiff.imwrite(os.path.join(dirs["Prep 1->2"], fname), preprocess_combined(img_raw))
+        
     return dirs
 
 def run_nellie_experiment(dataset_dir, combined_data_path, run_native_filter):
     files = sorted(glob.glob(os.path.join(dataset_dir, "*.TIF")))
     results, saved_masks = [], {}
-
+    
     for filepath in tqdm(files, desc=f"Evaluating", leave=False):
         filename = os.path.basename(filepath)
         gt_mask_path = os.path.join(combined_data_path, filename.replace(".TIF", "_seg.npy"))
         if not os.path.exists(gt_mask_path): continue
-
+        
         f_info = FileInfo(filepath=filepath)
         f_info.find_metadata()
         f_info.load_metadata()
@@ -142,88 +137,105 @@ def run_nellie_experiment(dataset_dir, combined_data_path, run_native_filter):
         f_info.axes = "ZYX" if img_array.ndim == 3 else "YX"
         f_info.good_axes, f_info.good_dims = True, True
         im_info = ImInfo(file_info=f_info)
-
+        
         if run_native_filter:
             frangi_filter = Filter(im_info=im_info, device="auto")
             frangi_filter.run()
         else:
             prep_path = im_info.create_output_path(pipeline_path="im_preprocessed", ext=".ome.tif")
             im_info.allocate_memory(output_path=prep_path, data=img_array, dtype=str(img_array.dtype))
-        
+            
         label_processor = Label(im_info=im_info, otsu_thresh_intensity=False, min_radius_um=10.0, device="auto")
         label_processor.run()
-
         pred_mask = np.squeeze(tiff.imread(im_info.pipeline_paths["im_instance_label"]))
         
         gt_data = np.load(gt_mask_path, allow_pickle=True)
         if gt_data.ndim == 0 and isinstance(gt_data.item(), dict):
             gt_data = gt_data.item().get("masks", gt_data.item())
         gt_to_compare = np.squeeze(gt_data)
-
-        # Strict IoU = 0.5 Evaluation
+        
         metrics = evaluate_predictions(pred_mask, gt_to_compare, iou_threshold=0.5)
         results.append(metrics)
         saved_masks[filename] = (pred_mask, metrics[4])
-
+        
     return results, saved_masks
 
-
 # =====================================================================
-# Plot Single Experiment (Triggers inside loop)
+# Plot Single Experiment 
 # =====================================================================
-def plot_single_experiment_visual(sample_filename, combined_data_path, data_src_dir, config_name, pred_mask, pred_count):
+def plot_single_experiment_visual(sample_filename, combined_data_path, data_src_dir, raw_data_dir, config_name, pred_mask, pred_count):
     logger.info(f"Generating visual plot for: {config_name}")
     
+    # --- Load Data ---
     gt_data = np.load(os.path.join(combined_data_path, sample_filename.replace(".TIF", "_seg.npy")), allow_pickle=True)
     if gt_data.ndim == 0 and isinstance(gt_data.item(), dict):
         gt_data = gt_data.item().get("masks", gt_data.item())
     gt_mask = np.squeeze(gt_data)
-    
     gt_count = len(np.unique(gt_mask)[np.unique(gt_mask) != 0])
-    gt_boundaries = find_boundaries(gt_mask, mode="inner")
     
     feed_img = tiff.imread(os.path.join(data_src_dir, sample_filename))
-    pred_boundaries = find_boundaries(pred_mask, mode="inner")
+    raw_img = tiff.imread(os.path.join(raw_data_dir, sample_filename))
+    
+    # --- Ground Truth High-Def Assets ---
+    gt_binary = gt_mask > 0
+    gt_boundaries = find_boundaries(gt_binary, mode="outer")
+    gt_thick_boundaries = binary_dilation(gt_boundaries, disk(1))
+    gt_masked_fill = np.ma.masked_where(~gt_binary, gt_binary)
+    gt_masked_edges = np.ma.masked_where(~gt_thick_boundaries, gt_thick_boundaries)
+
+    # --- Prediction High-Def Assets ---
+    pred_binary = pred_mask > 0
+    pred_boundaries = find_boundaries(pred_binary, mode="outer")
+    pred_thick_boundaries = binary_dilation(pred_boundaries, disk(1))
+    pred_masked_fill = np.ma.masked_where(~pred_binary, pred_binary)
+    pred_masked_edges = np.ma.masked_where(~pred_thick_boundaries, pred_thick_boundaries)
+    
     cmap_choice = "gray" if "Prep 1" in data_src_dir or "Raw" in data_src_dir else "magma"
-
     fig, axes = plt.subplots(1, 3, figsize=(18, 6))
-
-    # 1. Ground Truth Overlay
-    axes[0].imshow(feed_img, cmap=cmap_choice)
-    axes[0].imshow(np.dstack([gt_boundaries, np.zeros_like(gt_boundaries), np.zeros_like(gt_boundaries), gt_boundaries.astype(float)]))
-    axes[0].set_title(f"GT Overlay (Count: {gt_count})", fontsize=12, fontweight="bold")
+    
+    # 1. Ground Truth Overlay (on Raw Image)
+    axes[0].imshow(raw_img, cmap="gray")
+    axes[0].imshow(gt_masked_fill, cmap="summer", alpha=0.25) # Greenish overlay
+    axes[0].imshow(gt_masked_edges, cmap="Greens", alpha=0.9)
+    axes[0].set_title(f"GT Overlay on Raw (Count: {gt_count})", fontsize=12, fontweight="bold")
     axes[0].axis("off")
-
-    # 2. Input Image Alone
+    
+    # 2. Input Image Alone (The preprocessed image actually fed to Nellie)
     axes[1].imshow(feed_img, cmap=cmap_choice)
     axes[1].set_title(f"Data Fed into Nellie", fontsize=12, fontweight="bold")
     axes[1].axis("off")
-
-    # 3. Predicted Mask Overlay
-    axes[2].imshow(feed_img, cmap=cmap_choice)
-    overlay = np.zeros((*pred_mask.shape, 4))
-    overlay[pred_boundaries] = [1.0, 0.0, 0.0, 1.0]
-    axes[2].imshow(overlay)
-    axes[2].set_title(f"Predicted Mask (Count: {pred_count})", fontsize=12, fontweight="bold")
+    
+    # 3. Predicted Mask Overlay (on Raw Image)
+    axes[2].imshow(raw_img, cmap="gray")
+    axes[2].imshow(pred_masked_fill, cmap="autumn", alpha=0.25) # Reddish/Autumn overlay
+    axes[2].imshow(pred_masked_edges, cmap="Wistia", alpha=0.9)
+    axes[2].set_title(f"Predicted Mask on Raw (Count: {pred_count})", fontsize=12, fontweight="bold")
     axes[2].axis("off")
-
+    
     plt.suptitle(f"EXPERIMENT: {config_name}\nFile: {sample_filename}", fontsize=16, fontweight="bold")
     plt.tight_layout()
     plt.show()
 
-
 # =====================================================================
 # Main Execution Flow
 # =====================================================================
 def main():
+    n_samples = 10  # Default set to 10
+    
     logger.info("Downloading dataset...")
     dataset_dir = snapshot_download(repo_id="champ7/CEllDataW3", repo_type="dataset", allow_patterns="combined_data/*")
     combined_data_path = os.path.join(dataset_dir, "combined_data")
+    
+    # Randomly sample N files to drastically speed up testing
     raw_tif_files = sorted(glob.glob(os.path.join(combined_data_path, "*.TIF")))
+    if n_samples and n_samples < len(raw_tif_files):
+        random.seed(42) # Ensure the same random subset is used across runs
+        raw_tif_files = random.sample(raw_tif_files, n_samples)
+        logger.info(f"Subsampled dataset to {n_samples} random images.")
+        
     sample_filename = os.path.basename(raw_tif_files[0]) # File to plot for each step
-
     dirs = prepare_datasets(raw_tif_files)
-
+    
     configs = [
         {"name": "Prep 1 Only (Bypassing Nellie Filter)",       "data_src": "Prep 1",    "run_filter": False},
         {"name": "Prep 2 Only (Bypassing Nellie Filter)",       "data_src": "Prep 2",    "run_filter": False},
@@ -233,10 +245,10 @@ def main():
         {"name": "Prep 2 + Nellie Native Filter",               "data_src": "Prep 2",    "run_filter": True},
         {"name": "Prep 1->2 + Nellie Native Filter",            "data_src": "Prep 1->2", "run_filter": True},
     ]
-
+    
     aggregate_table = []
     headers = ["Strategy Configuration", "Total GT Cells", "Total Pred Cells", "Mean Pixel IoU", "Mean Precision", "Mean Recall", "Mean F1"]
-
+    
     for config in configs:
         logger.info(f"Running strategy: {config['name']}")
         
@@ -248,7 +260,7 @@ def main():
         
         ious, precs, recs, f1s = [m[0] for m in metrics], [m[1] for m in metrics], [m[2] for m in metrics], [m[3] for m in metrics]
         total_pred, total_gt = sum([m[4] for m in metrics]), sum([m[5] for m in metrics])
-
+        
         current_result_row = [
             config["name"], 
             total_gt, total_pred, 
@@ -262,13 +274,14 @@ def main():
         print(f"="*80)
         print(tabulate([current_result_row], headers=headers, tablefmt="fancy_grid"))
         print("\n")
-
+        
         # 2. Trigger the plot immediately
         pred_mask_sample, pred_count_sample = saved_masks[sample_filename]
         plot_single_experiment_visual(
             sample_filename=sample_filename,
             combined_data_path=combined_data_path,
             data_src_dir=dirs[config["data_src"]],
+            raw_data_dir=dirs["Raw"],
             config_name=config["name"],
             pred_mask=pred_mask_sample,
             pred_count=pred_count_sample