Create stardistw2.py

stardistw2.py

@@ -0,0 +1,244 @@
+import os
+import glob
+import cv2
+import numpy as np
+import tifffile
+import pandas as pd
+import matplotlib.subplots as plt_subplots
+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 csbdeep.utils import normalize
+from stardist.models import Config2D, StarDist2D
+from stardist import fill_label_holes, calculate_extents, random_label_cmap
+from stardist.matching import matching
+
+# --- 1. Data Loading & Preparation ---
+def load_and_split_dataset(repo_id="champ7/celldatamag", test_size=10):
+    print(f"Downloading dataset from Hugging Face: {repo_id}...")
+    local_dir = snapshot_download(repo_id=repo_id, repo_type="dataset")
+    
+    # Target specifically _w2.TIF based on the dataset structure
+    search_path = os.path.join(local_dir, "**", "*_w2.TIF")
+    image_paths = sorted(glob.glob(search_path, recursive=True))
+    
+    X, Y = [], []
+    for img_path in image_paths:
+        # Map _w2.TIF to _w3_seg.npy based on the provided screenshot
+        mask_path = img_path.replace("_w2.TIF", "_w3_seg.npy")
+        
+        if os.path.exists(mask_path):
+            img = tifffile.imread(img_path)
+            mask = np.load(mask_path, allow_pickle=True)
+            
+            if mask.ndim == 0 and isinstance(mask.item(), dict):
+                mask = mask.item().get('masks', mask.item())
+            
+            img = np.squeeze(img)
+            mask = np.squeeze(mask)
+                
+            mask = fill_label_holes(mask.astype(np.int32))
+            X.append(img)
+            Y.append(mask)
+
+    print(f"Loaded {len(X)} valid 2D image-mask pairs.")
+    
+    if len(X) <= test_size:
+        raise ValueError("Not enough images to create a split. Reduce test_size or check the dataset.")
+
+    X_test, Y_test = X[:test_size], Y[:test_size]
+    X_train, Y_train = X[test_size:], Y[test_size:]
+    
+    return X_train, Y_train, X_test, Y_test
+
+# --- 2. Preprocessing Definitions ---
+def process_w2(image_obj):
+    """ W2: Organelles (Red Fluorescence) -> CLAHE & Top-Hat preprocessing """
+    if image_obj.ndim == 3: 
+        image_obj = image_obj[0]
+    p1, p99 = np.percentile(image_obj, (1, 99.9))
+    robust_img = exposure.rescale_intensity(image_obj, in_range=(p1, p99), out_range=(0, 255)).astype(np.uint8)
+    clahe = cv2.createCLAHE(clipLimit=5.0, tileGridSize=(8, 8))
+    clahe_img = clahe.apply(robust_img)
+    return clahe_img
+
+def preprocess_2(img):
+    """ Standard StarDist Percentile Normalization """
+    return normalize(img, 1, 99.8, axis=(0,1))
+
+def preprocess_none(img):
+    """ Basic Min-Max scaling to prevent float crashes on raw inputs """
+    return (img - np.min(img)) / (np.ptp(img) + 1e-8)
+
+def apply_preprocessing(img, mode="none", to_rgb=False):
+    """ Router for preprocessing modes """
+    if mode == "prep1":
+        out = process_w2(img)
+    elif mode == "prep2":
+        out = preprocess_2(img)
+    elif mode == "prep1_then_prep2":
+        out = process_w2(img)
+        out = preprocess_2(out)
+    else:
+        out = preprocess_none(img)
+        
+    out = out.astype(np.float32)
+
+    if to_rgb and out.ndim == 2:
+        out = np.stack((out,) * 3, axis=-1)
+        
+    return out
+
+# --- 3. Visualization Function ---
+def plot_test_results(X_test, Y_test, Y_pred, X_input, experiment_name, mode_label, num_images=2):
+    """Plots 3 columns: Original+GT, Model Input, Model Input+Pred"""
+    print(f"\nGenerating plots for: {experiment_name}")
+    lbl_cmap = random_label_cmap()
+    
+    for i in range(min(num_images, len(X_test))):
+        fig, axes = plt.subplots(1, 3, figsize=(18, 6))
+        
+        # Column 1: Original Image + GT
+        axes[0].imshow(X_test[i], cmap='gray')
+        axes[0].imshow(Y_test[i], cmap=lbl_cmap, alpha=0.5)
+        axes[0].set_title(f"Image {i+1}: Original + Ground Truth")
+        axes[0].axis('off')
+        
+        # Column 2: The actual input fed to the model
+        disp_img = X_input[i]
+        if disp_img.ndim == 3: 
+            disp_img = disp_img[..., 0] 
+            
+        axes[1].imshow(disp_img, cmap='gray')
+        axes[1].set_title(f"Image {i+1}: Model Input ({mode_label})")
+        axes[1].axis('off')
+        
+        # Column 3: Model Input + Prediction
+        axes[2].imshow(disp_img, cmap='gray')
+        axes[2].imshow(Y_pred[i], cmap=lbl_cmap, alpha=0.5)
+        axes[2].set_title(f"Image {i+1}: Model Input + Prediction")
+        axes[2].axis('off')
+        
+        plt.tight_layout()
+        plt.show()
+
+# --- 4. Training Function ---
+def train_model(X_train, Y_train, model_name, prep_mode="none"):
+    print(f"\n--- Training Custom Model ONCE: {model_name} (Training Data Preprocessing: {prep_mode}) ---")
+    
+    X_trn_proc = [apply_preprocessing(x, mode=prep_mode, to_rgb=False) for x in X_train]
+    
+    rng = np.random.RandomState(42)
+    ind = rng.permutation(len(X_trn_proc))
+    n_val = max(1, int(round(0.15 * len(ind))))
+    ind_train, ind_val = ind[:-n_val], ind[-n_val:]
+    
+    X_trn = [X_trn_proc[i] for i in ind_train]
+    Y_trn = [Y_train[i] for i in ind_train]
+    X_val = [X_trn_proc[i] for i in ind_val]
+    Y_val = [Y_train[i] for i in ind_val]
+    
+    conf = Config2D(
+        n_rays=32,
+        grid=(2,2),
+        n_channel_in=1,
+        train_patch_size=(256,256),
+        train_batch_size=4,
+    )
+    model = StarDist2D(conf, name=model_name, basedir='models')
+    
+    median_size = calculate_extents(list(Y_train), np.median)
+    fov = np.array(model._axes_tile_overlap('YX'))
+    if any(median_size > fov):
+        print("WARNING: Median object size is larger than the field of view.")
+    
+    model.train(X_trn, Y_trn, validation_data=(X_val, Y_val), epochs=30, steps_per_epoch=50)
+    model.optimize_thresholds(X_val, Y_val)
+    
+    return model
+
+# --- 5. Evaluation Loop ---
+def run_evaluation(model, X_test, Y_test, experiment_name, prep_mode, is_pretrained_he=False):
+    print(f"\nEvaluating: {experiment_name}")
+    metrics = []
+    
+    Y_pred_list = []
+    X_input_list = []
+    
+    for i, (img, gt) in enumerate(zip(X_test, Y_test)):
+        img_input = apply_preprocessing(img, mode=prep_mode, to_rgb=is_pretrained_he)
+        X_input_list.append(img_input)
+        
+        pred_mask, _ = model.predict_instances(img_input, prob_thresh=0.5, nms_thresh=0.3)
+        Y_pred_list.append(pred_mask)
+        
+        # IoU Threshold 0.0 for matching calculation
+        res = matching(gt, pred_mask, thresh=0.0)
+        
+        metrics.append({
+            "Image": i+1, 
+            "F1": res.f1, 
+            "IoU": res.mean_matched_score 
+        })
+    
+    df = pd.DataFrame(metrics)
+    mean_f1 = df['F1'].mean()
+    mean_iou = df['IoU'].mean()
+    
+    print(f"Results for {experiment_name}:")
+    print(f"Mean F1-Score: {mean_f1:.4f}")
+    print(f"Mean IoU:      {mean_iou:.4f}\n")
+    
+    plot_test_results(X_test, Y_test, Y_pred_list, X_input_list, experiment_name, prep_mode, num_images=2)
+    
+    return mean_f1, mean_iou
+
+# --- Main Execution Block ---
+if __name__ == "__main__":
+    np.random.seed(42) 
+    
+    # 1. Load Data with updated repo
+    X_train, Y_train, X_test, Y_test = load_and_split_dataset(repo_id="champ7/celldatamag", test_size=10)
+    
+    all_results = []
+    
+    # 2. Preload the versatile H&E model
+    print("\nLoading pre-trained '2D_versatile_he' model...")
+    pretrained_model = StarDist2D.from_pretrained('2D_versatile_he')
+
+    # 3. Train the custom model EXACTLY ONCE (using 'none' as the baseline training preprocessing)
+    custom_trained_model = train_model(X_train, Y_train, model_name="single_custom_model", prep_mode="none")
+
+    # 4. Define all 8 experiments
+    experiments = [
+        # (Experiment Name, Preprocessing Mode, Model to Evaluate, is_pretrained_flag)
+        ("Exp 1: Pretrained, No Preprocessing", "none", pretrained_model, True),
+        ("Exp 2: Pretrained, Preprocessing 1", "prep1", pretrained_model, True),
+        ("Exp 3: Pretrained, Preprocessing 2", "prep2", pretrained_model, True),
+        ("Exp 4: Pretrained, Prep 1 -> Prep 2", "prep1_then_prep2", pretrained_model, True),
+        
+        ("Exp 5: Trained, No Preprocessing", "none", custom_trained_model, False),
+        ("Exp 6: Trained, Preprocessing 1", "prep1", custom_trained_model, False),
+        ("Exp 7: Trained, Preprocessing 2", "prep2", custom_trained_model, False),
+        ("Exp 8: Trained, Prep 1 -> Prep 2", "prep1_then_prep2", custom_trained_model, False),
+    ]
+
+    # 5. Run through all inference/evaluation states
+    for exp_name, prep_mode, model_to_eval, is_pretrained in experiments:
+        f1, iou = run_evaluation(model_to_eval, X_test, Y_test, exp_name, prep_mode, is_pretrained_he=is_pretrained)
+        all_results.append({
+            "Experiment": exp_name,
+            "Mean F1-Score": round(f1, 4),
+            "Mean IoU": round(iou, 4)
+        })
+
+    # --- 6. Final Summary Table ---
+    print("\n" + "="*70)
+    print("FINAL METRICS SUMMARY (IoU Threshold = 0.0)".center(70))
+    print("="*70)
+    
+    results_df = pd.DataFrame(all_results)
+    print(results_df.to_string(index=False, justify='left'))
+    print("="*70)