Create omnipose.py

omnipose.py

@@ -0,0 +1,266 @@
+import os
+import glob
+import random
+import numpy as np
+import pandas as pd
+import tifffile as tiff
+import matplotlib.pyplot as plt
+from scipy.ndimage import median_filter
+from skimage.filters import gaussian
+from skimage import exposure
+from huggingface_hub import snapshot_download
+from cellpose_omni import models, metrics, core
+
+# ==========================================
+# 1. Custom Preprocessing (Prep 1)
+# ==========================================
+def preprocess_w3(input_data):
+    """ W3: Cells (Brightfield) -> Blur, Median, Contrast """
+    if isinstance(input_data, str):
+        img_raw = tiff.imread(input_data).astype(np.float32)
+    else:
+        img_raw = input_data.astype(np.float32)
+        
+    img_min, img_max = np.min(img_raw), np.max(img_raw)
+    img_8bit = ((img_raw - 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: 
+        final_img = img_filtered
+    else:
+        final_img = exposure.rescale_intensity(img_filtered, in_range=(p_low, p_high), out_range=(0, 255)).astype(np.uint8)
+        
+    # Standardize to 0.0 - 1.0 floats
+    return (final_img / 255.0).astype(np.float32)
+
+# ==========================================
+# 2. Data Preparation
+# ==========================================
+def prepare_data(repo_id="champ7/CEllDataW3"):
+    print(f"Downloading dataset {repo_id} from Hugging Face...")
+    dataset_path = snapshot_download(repo_id=repo_id, repo_type="dataset")
+    data_dir = os.path.join(dataset_path, "combined_data")
+    
+    image_files = sorted(glob.glob(os.path.join(data_dir, "*.TIF")))
+    seg_files = sorted(glob.glob(os.path.join(data_dir, "*_seg.npy")))
+    
+    images, masks = [], []
+    
+    for img_path in image_files:
+        base_name = img_path.replace(".TIF", "")
+        seg_path = base_name + "_seg.npy"
+        
+        if seg_path in seg_files:
+            img = tiff.imread(img_path).astype(np.float32)
+            images.append(img)
+            
+            seg_data = np.load(seg_path, allow_pickle=True)
+            if seg_data.shape == ():
+                seg_dict = seg_data.item()
+                mask_raw = seg_dict['masks']
+            else:
+                mask_raw = seg_data
+                
+            # Convert to standard 32-bit integer labels
+            masks.append(mask_raw.astype(np.int32))
+            
+    X_train, y_train = images, masks
+    
+    combined = list(zip(images, masks))
+    random.seed(42)
+    sample_10 = random.sample(combined, min(10, len(combined)))
+    X_test, y_test = zip(*sample_10)
+    
+    return list(X_train), list(X_test), list(y_train), list(y_test)
+
+# ==========================================
+# 3. Metrics & Utilities
+# ==========================================
+def compute_metrics(masks_true, masks_pred, threshold=0.0):
+    ap, tp, fp, fn = metrics.average_precision(masks_true, masks_pred, threshold=[threshold])
+    
+    tp_sum = np.sum(tp)
+    fp_sum = np.sum(fp)
+    fn_sum = np.sum(fn)
+    
+    f1_score = 0.0 if (2 * tp_sum + fp_sum + fn_sum) == 0 else (2 * tp_sum) / (2 * tp_sum + fp_sum + fn_sum)
+    mean_ap = np.mean(ap)
+    
+    return f1_score, mean_ap
+
+def to_numpy(data):
+    if hasattr(data, 'cpu'):
+        return data.cpu().detach().numpy()
+    elif isinstance(data, list):
+        return [to_numpy(d) for d in data]
+    elif isinstance(data, tuple):
+        return tuple(to_numpy(d) for d in data)
+    elif isinstance(data, dict):
+        return {k: to_numpy(v) for k, v in data.items()}
+    elif isinstance(data, np.ndarray) and data.dtype == object:
+        return np.array([to_numpy(d) for d in data])
+    return data
+
+# ==========================================
+# 4. Visualization Helper (4 Panels)
+# ==========================================
+def visualize_results(raw_img, fed_img, y_true, y_pred, config_name, model_type, img_idx):
+    fig, axes = plt.subplots(1, 4, figsize=(24, 6))
+    fig.suptitle(f"{config_name} | {model_type} | Image {img_idx}", fontsize=16, fontweight='bold')
+
+    axes[0].imshow(raw_img, cmap='gray')
+    axes[0].set_title("Original Image", fontsize=14)
+    axes[0].axis('off')
+
+    axes[1].imshow(raw_img, cmap='gray')
+    mask_true_display = np.ma.masked_where(y_true == 0, y_true)
+    axes[1].imshow(mask_true_display, cmap='nipy_spectral', alpha=0.5, interpolation='none')
+    axes[1].set_title("Original Image + GT Mask", fontsize=14)
+    axes[1].axis('off')
+
+    axes[2].imshow(fed_img, cmap='gray')
+    axes[2].set_title("Preprocessed Image (Fed to Model)", fontsize=14)
+    axes[2].axis('off')
+
+    axes[3].imshow(raw_img, cmap='gray')
+    mask_pred_display = np.ma.masked_where(y_pred == 0, y_pred)
+    axes[3].imshow(mask_pred_display, cmap='nipy_spectral', alpha=0.5, interpolation='none')
+    axes[3].set_title("Original Image + Predicted Mask", fontsize=14)
+    axes[3].axis('off')
+
+    plt.tight_layout()
+    plt.show()
+
+# ==========================================
+# 5. Pipeline Execution
+# ==========================================
+def run_pipeline():
+    # Use GPU for fast inference/evaluation
+    use_gpu_eval = core.use_gpu()
+    channels = [0, 0]
+    base_model_name = 'cyto2'
+    
+    X_train_raw, X_test_raw, y_train, y_test = prepare_data()
+    
+    X_train_p1 = [preprocess_w3(img) for img in X_train_raw]
+    X_test_p1 = [preprocess_w3(img) for img in X_test_raw]
+
+    configs = [
+        {
+            "name": "Prep 1", 
+            "train_data": X_train_p1, "test_data": X_test_p1,
+            "train_kwargs": {"normalize": False, "rescale": None},
+            "eval_kwargs": {"normalize": False, "flow_threshold": 0, "mask_threshold": -1, "rescale": None}
+        },
+        {
+            "name": "Prep 2", 
+            "train_data": X_train_raw, "test_data": X_test_raw,
+            "train_kwargs": {"normalize": True, "rescale": None},
+            "eval_kwargs": {"normalize": True, "flow_threshold": 0, "mask_threshold": -1, "rescale": None}
+        },
+        {
+            "name": "Prep 1 + Prep 2", 
+            "train_data": X_train_p1, "test_data": X_test_p1,
+            "train_kwargs": {"normalize": True, "rescale": None},
+            "eval_kwargs": {"normalize": True, "flow_threshold": 0, "mask_threshold": -1, "rescale": None}
+        }
+    ]
+
+    results = []
+
+    for cfg in configs:
+        print(f"\n--- Running Configuration: {cfg['name']} ---")
+        
+        # ------------------------------------------
+        # A. Pretrained Model Evaluation (Runs on GPU if available)
+        # ------------------------------------------
+        pretrained_model = models.CellposeModel(
+            gpu=use_gpu_eval, 
+            model_type=base_model_name, 
+            omni=True, 
+            nchan=2, 
+            nclasses=2
+        )
+        preds_pre, _, _ = pretrained_model.eval(cfg["test_data"], channels=channels, **cfg["eval_kwargs"])
+        
+        preds_pre = to_numpy(preds_pre)
+        f1_pre, iou_pre = compute_metrics(y_test, preds_pre, threshold=0.0)
+        
+        results.append({
+            "Configuration": f"{cfg['name']} + Pretrained + Prediction",
+            "F1 Score": f"{f1_pre:.4f}",
+            "Mean IoU (AP)": f"{iou_pre:.4f}"
+        })
+        
+        for viz_idx in range(min(2, len(cfg["test_data"]))):
+            visualize_results(
+                raw_img=X_test_raw[viz_idx], fed_img=cfg["test_data"][viz_idx], 
+                y_true=y_test[viz_idx], y_pred=preds_pre[viz_idx],
+                config_name=cfg["name"], model_type="Pretrained Model", img_idx=viz_idx + 1
+            )
+        
+        # ------------------------------------------
+        # B. Fine-tuning Model (gpu=False BYPASSES CUDA LOSS ASSERTIONS COMPLETELY)
+        # ------------------------------------------
+        print(f"Fine-tuning {base_model_name} on CPU to guarantee zero CUDA assertion crashes...")
+        finetune_model = models.CellposeModel(
+            #gpu=False,  # <--- THIS REMOVES THE ASSERTION PROBLEM ENTIRELY
+            model_type=base_model_name, 
+            omni=True,
+            nchan=2,
+            nclasses=2
+        )
+        
+        model_path = finetune_model.train(
+            cfg["train_data"], y_train, 
+            train_links=[None] * len(y_train), 
+            channels=channels, 
+            save_path=f"./omnipose_{cfg['name'].replace(' ', '_')}", 
+            n_epochs=150, 
+            learning_rate=0.1,
+            batch_size=2,
+            **cfg["train_kwargs"]
+        )
+        
+        # ------------------------------------------
+        # C. Evaluate Fine-Tuned Model (Evaluates back on GPU)
+        # ------------------------------------------
+        custom_model = models.CellposeModel(
+            gpu=use_gpu_eval, 
+            pretrained_model=model_path, 
+            model_type=base_model_name, 
+            omni=True,
+            nchan=2,
+            nclasses=2
+        )
+        preds_fine, _, _ = custom_model.eval(cfg["test_data"], channels=channels, **cfg["eval_kwargs"])
+        
+        preds_fine = to_numpy(preds_fine)
+        f1_fine, iou_fine = compute_metrics(y_test, preds_fine, threshold=0.0)
+        
+        results.append({
+            "Configuration": f"{cfg['name']} + Training + Prediction",
+            "F1 Score": f"{f1_fine:.4f}",
+            "Mean IoU (AP)": f"{iou_fine:.4f}"
+        })
+        
+        for viz_idx in range(min(2, len(cfg["test_data"]))):
+            visualize_results(
+                raw_img=X_test_raw[viz_idx], fed_img=cfg["test_data"][viz_idx], 
+                y_true=y_test[viz_idx], y_pred=preds_fine[viz_idx],
+                config_name=cfg["name"], model_type="Fine-tuned Model", img_idx=viz_idx + 1
+            )
+
+    # ==========================================
+    # 6. Final Output Compilation
+    # ==========================================
+    print("\n================ FINAL METRICS TABLE ================")
+    results_df = pd.DataFrame(results)
+    print(results_df.to_markdown(index=False))
+
+if __name__ == "__main__":
+    run_pipeline()