Create nellie.py

nellie.py

@@ -0,0 +1,284 @@
+import os
+import glob
+import logging
+import numpy as np
+import matplotlib.pyplot as plt
+import tifffile as tiff
+from scipy.sparse import csr_matrix
+from tabulate import tabulate
+from tqdm import tqdm
+from huggingface_hub import snapshot_download
+
+from nellie.im_info.verifier import FileInfo, ImInfo
+from nellie.segmentation.filtering import Filter
+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.segmentation import find_boundaries
+from skimage.util import img_as_ubyte
+
+# Logging setup
+logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s", datefmt="%H:%M:%S")
+logger = logging.getLogger(__name__)
+
+# =====================================================================
+# Preprocessing Logic
+# =====================================================================
+def preprocess_method_1(img_raw: np.ndarray) -> np.ndarray:
+    img_float = img_raw.astype(np.float32)
+    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)
+
+def preprocess_method_2(img_input: np.ndarray) -> np.ndarray:
+    img_min, img_max = img_input.min(), img_input.max()
+    img_normalized = (img_input - img_min) / (img_max - img_min if img_max > img_min else 1.0)
+    img_8bit = img_as_ubyte(img_normalized)
+    nellie_input = rank.entropy(img_8bit, disk(25))
+    edge_map = sobel(img_normalized)
+    texture_features = nellie_input * (1.0 + edge_map * 5.0)
+    tex_min, tex_max = texture_features.min(), texture_features.max()
+    tex_normalized = (texture_features - tex_min) / (tex_max - tex_min if tex_max > tex_min else 1.0)
+    return (tex_normalized * 255).astype(np.uint8)
+
+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)
+# =====================================================================
+def evaluate_predictions(pred_mask: np.ndarray, gt_mask: np.ndarray, iou_threshold: float = 0.5):
+    """Vectorized instance-level matching using strict IoU >= 0.5"""
+    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]
+    pred_vec = np.array([np.bincount(pred_mask.ravel())[id_] for id_ in pred_labels])[None, :]
+    
+    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
+    
+    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
+# =====================================================================
+def prepare_datasets(raw_tif_files, base_dir="generated_datasets"):
+    dirs = {
+        "Raw": os.path.join(base_dir, "dataset_raw"),
+        "Prep 1": os.path.join(base_dir, "dataset_prep1"),
+        "Prep 2": os.path.join(base_dir, "dataset_prep2"),
+        "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)):
+            continue
+            
+        img_raw = np.squeeze(tiff.imread(filepath))
+        tiff.imwrite(os.path.join(dirs["Raw"], fname), img_raw)
+        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()
+        
+        img_array = np.squeeze(tiff.imread(filepath))
+        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)
+# =====================================================================
+def plot_single_experiment_visual(sample_filename, combined_data_path, data_src_dir, config_name, pred_mask, pred_count):
+    logger.info(f"Generating visual plot for: {config_name}")
+    
+    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")
+    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")
+    axes[0].axis("off")
+
+    # 2. Input Image Alone
+    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")
+    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():
+    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")
+    raw_tif_files = sorted(glob.glob(os.path.join(combined_data_path, "*.TIF")))
+    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},
+        {"name": "Prep 1->2 Only (Bypassing Nellie Filter)",    "data_src": "Prep 1->2", "run_filter": False},
+        {"name": "Nellie Native Only (Raw -> Nellie Filter)",   "data_src": "Raw",       "run_filter": True},
+        {"name": "Prep 1 + Nellie Native Filter",               "data_src": "Prep 1",    "run_filter": True},
+        {"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']}")
+        
+        metrics, saved_masks = run_nellie_experiment(
+            dataset_dir=dirs[config["data_src"]], 
+            combined_data_path=combined_data_path, 
+            run_native_filter=config["run_filter"]
+        )
+        
+        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, 
+            f"{np.mean(ious):.4f}", f"{np.mean(precs):.4f}", f"{np.mean(recs):.4f}", f"{np.mean(f1s):.4f}"
+        ]
+        aggregate_table.append(current_result_row)
+        
+        # 1. Print immediate feedback for this run
+        print(f"\n" + "="*80)
+        print(f" EXPERIMENT RESULTS: {config['name']}")
+        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"]],
+            config_name=config["name"],
+            pred_mask=pred_mask_sample,
+            pred_count=pred_count_sample
+        )
+
+    # Print Final Summary Table
+    print("\n" + "=" * 115)
+    print(" " * 35 + f"FINAL 7-WAY EVALUATION REPORT (N={len(raw_tif_files)} images)")
+    print("=" * 115)
+    print(tabulate(aggregate_table, headers=headers, tablefmt="fancy_grid"))
+
+if __name__ == "__main__":
+    main()