Create microsam.py

microsam.py

@@ -0,0 +1,285 @@
+!pip install numpy tifffile scikit-image matplotlib huggingface_hub scikit-learn micro_sam
+!conda install -c conda-forge micro_sam
+
+import os
+import glob
+import time
+import numpy as np
+import tifffile as tiff
+import matplotlib.pyplot as plt
+from skimage.filters import gaussian, median
+from skimage.morphology import square
+from skimage import exposure
+from huggingface_hub import snapshot_download
+from sklearn.metrics import f1_score
+
+# --- micro_sam imports ---
+import micro_sam.training as sam_training
+from micro_sam.util import export_custom_sam_model
+from micro_sam.automatic_segmentation import (
+    get_predictor_and_segmenter, 
+    automatic_instance_segmentation
+)
+
+
+# ==========================================
+# 1. PREPROCESSING FUNCTIONS
+# ==========================================
+
+def preprocess_1(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)
+    if img_max > img_min:
+        img_8bit = ((img_raw - img_min) / (img_max - img_min) * 255.0).astype(np.uint8) 
+    else:
+        img_8bit = 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(img_blur, footprint=square(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.float32)
+        
+    # Output float in range [0, 255] for micro_sam dataloader compatibility
+    return final_img.astype(np.float32)
+
+def preprocess_2(input_data):
+    """ Percentile Clipping (1st, 99.9th), Normalization, Standardization -> Scaled to [0, 255] """
+    if isinstance(input_data, str):
+        img = tiff.imread(input_data).astype(np.float32)
+    else:
+        img = input_data.astype(np.float32)
+        
+    # Percentile clipping
+    p_low, p_high = np.percentile(img, (1.0, 99.9))
+    img_clipped = np.clip(img, p_low, p_high)
+    
+    # Scale to [0, 1]
+    denom = (p_high - p_low) if (p_high - p_low) > 0 else 1.0
+    img_scaled = (img_clipped - p_low) / denom
+    
+    # Zero-mean, unit-variance standardization
+    mean, std = np.mean(img_scaled), np.std(img_scaled)
+    img_standardized = (img_scaled - mean) / (std + 1e-8)
+    
+    # Scale exactly to [0, 255] to bypass micro_sam ValueError
+    s_min, s_max = np.min(img_standardized), np.max(img_standardized)
+    if s_max > s_min:
+        img_final = ((img_standardized - s_min) / (s_max - s_min) * 255.0)
+    else:
+        img_final = np.zeros_like(img_standardized)
+        
+    return img_final.astype(np.float32)
+
+def preprocess_1_then_2(input_data):
+    """ Sequential application of Preprocessing 1 followed by Preprocessing 2 """
+    out_1 = preprocess_1(input_data)
+    out_2 = preprocess_2(out_1)
+    return out_2
+
+
+# ==========================================
+# 2. METRICS & PLOTTING
+# ==========================================
+
+def calculate_metrics(y_true, y_pred):
+    """ Calculate IoU and F1 score for binary masks. """
+    y_true_bin = (y_true > 0).astype(np.uint8)
+    y_pred_bin = (y_pred > 0).astype(np.uint8)
+    
+    intersection = np.logical_and(y_true_bin, y_pred_bin).sum()
+    union = np.logical_or(y_true_bin, y_pred_bin).sum()
+    
+    iou = intersection / union if union > 0 else 0.0
+    f1 = f1_score(y_true_bin.flatten(), y_pred_bin.flatten())
+    
+    return iou, f1
+
+def plot_results(title, actual_img, preprocessed_img, gt_mask, pred_mask, output_filename):
+    """ Generates 1x3 subplot and displays it instantly for review. """
+    fig, axes = plt.subplots(1, 3, figsize=(18, 6))
+    fig.suptitle(title, fontsize=16)
+    
+    # Image 1: Actual Image with GT Mask
+    axes[0].imshow(actual_img, cmap='gray')
+    axes[0].imshow(gt_mask > 0, cmap='Reds', alpha=0.3)
+    axes[0].set_title("Actual Image + GT Mask")
+    axes[0].axis('off')
+    
+    # Image 2: Preprocessed Image (Input to model)
+    axes[1].imshow(preprocessed_img, cmap='gray')
+    axes[1].set_title("Preprocessed Input")
+    axes[1].axis('off')
+    
+    # Image 3: Actual Image with Predicted Mask
+    axes[2].imshow(actual_img, cmap='gray')
+    axes[2].imshow(pred_mask > 0, cmap='Blues', alpha=0.3)
+    axes[2].set_title("Actual Image + Predicted Mask")
+    axes[2].axis('off')
+    
+    plt.tight_layout()
+    plt.savefig(output_filename)
+    
+    # Show the plot immediately so you can monitor direction
+    plt.show(block=False)
+    plt.pause(2) # Pauses briefly to render the UI before continuing
+    plt.close()
+
+
+# ==========================================
+# 3. PIPELINE & DATA HANDLING
+# ==========================================
+
+def download_data():
+    print("Downloading dataset 'champ7/CEllDataW3'...")
+    local_dir = snapshot_download(repo_id="champ7/CEllDataW3", repo_type="dataset")
+    return os.path.join(local_dir, "combined_data")
+
+def get_file_pairs(data_dir):
+    tif_files = sorted(glob.glob(os.path.join(data_dir, "*.TIF")))
+    pairs = []
+    for tif in tif_files:
+        mask_path = tif.replace(".TIF", "_seg.npy")
+        if os.path.exists(mask_path):
+            pairs.append((tif, mask_path))
+    return pairs
+
+def prepare_training_data(pairs, prep_func, prep_name):
+    """ Prepares data for the subset pairs. """
+    prep_dir = f"./preprocessed_train_data_{prep_name}"
+    os.makedirs(prep_dir, exist_ok=True)
+    
+    for img_path, mask_path in pairs:
+        base_name = os.path.basename(img_path)
+        img_prep = prep_func(img_path)
+        tiff.imwrite(os.path.join(prep_dir, base_name), img_prep)
+        
+        mask = np.load(mask_path)
+        tiff.imwrite(os.path.join(prep_dir, base_name.replace(".TIF", "_seg.tif")), mask.astype(np.int32))
+        
+    return prep_dir
+
+def train_model(train_dir, model_name):
+    print(f"\n--- Training {model_name} (Single Execution) ---")
+    
+    train_loader = sam_training.default_sam_loader(
+        raw_paths=train_dir,
+        raw_key="*.TIF",
+        label_paths=train_dir,
+        label_key="*_seg.tif",
+        patch_shape=(1, 512, 512),
+        batch_size=2,
+        with_segmentation_decoder=True,
+        is_train=True,
+        num_workers=2,
+        shuffle=True
+    )
+    
+    sam_training.train_sam(
+        name=model_name,
+        model_type="vit_b",
+        train_loader=train_loader,
+        val_loader=train_loader,
+        n_epochs=10, 
+        n_objects_per_batch=5,
+        with_segmentation_decoder=True,
+        device="cuda"
+    )
+    
+    checkpoint_path = os.path.join("checkpoints", model_name, "best.pt")
+    export_path = f"./{model_name}_exported.pth"
+    export_custom_sam_model(checkpoint_path=checkpoint_path, model_type="vit_b", save_path=export_path)
+    return export_path
+
+
+def run_inference_and_evaluate(pairs, prep_func, prep_name, model_type, model_checkpoint, is_pretrained=True):
+    model_status = "Pretrained" if is_pretrained else "Trained"
+    print(f"\nEvaluating {prep_name} with {model_status} Model...")
+    
+    predictor, segmenter = get_predictor_and_segmenter(
+        model_type=model_type,
+        checkpoint=model_checkpoint if not is_pretrained else None,
+    )
+    
+    total_iou, total_f1 = 0, 0
+    
+    for idx, (img_path, mask_path) in enumerate(pairs):
+        raw_img = tiff.imread(img_path)
+        gt_mask = np.load(mask_path)
+        
+        input_img = prep_func(raw_img)
+        pred_mask = automatic_instance_segmentation(predictor, segmenter, input_img)
+        
+        iou, f1 = calculate_metrics(gt_mask, pred_mask)
+        total_iou += iou
+        total_f1 += f1
+        
+        # Plot 2 images per combination as requested
+        if idx < 2:
+            plot_name = f"{prep_name}_{model_status}_sample_{idx+1}.png"
+            title = f"{prep_name} | {model_status} | IoU: {iou:.3f} | F1: {f1:.3f}"
+            plot_results(title, raw_img, input_img, gt_mask, pred_mask, plot_name)
+            
+    avg_iou = total_iou / len(pairs)
+    avg_f1 = total_f1 / len(pairs)
+    print(f"Finished {model_status} | {prep_name} -> F1: {avg_f1:.4f}, IoU: {avg_iou:.4f}")
+    return avg_iou, avg_f1
+
+
+# ==========================================
+# 4. MAIN EXECUTION
+# ==========================================
+
+if __name__ == "__main__":
+    data_dir = download_data()
+    
+    # CONSTRAINT MET: Only take 10 images from the dataset
+    file_pairs = get_file_pairs(data_dir)[:10]
+    print(f"Using a subset of {len(file_pairs)} image-mask pairs for faster processing.")
+
+    combinations = [
+        ("Preprocessing 1", preprocess_1, "prep1"),
+        ("Preprocessing 2", preprocess_2, "prep2"),
+        ("Preprocessing 1 + 2", preprocess_1_then_2, "prep12")
+    ]
+    
+    results = []
+
+    # 1. Evaluate Pretrained Model (3 Iterations)
+    for prep_title, prep_func, prep_tag in combinations:
+        iou, f1 = run_inference_and_evaluate(
+            file_pairs, prep_func, prep_title, 
+            model_type="vit_b", model_checkpoint=None, is_pretrained=True
+        )
+        results.append({"Pipeline": f"{prep_title} -> Pretrained Model", "F1": f1, "IoU@0.0": iou})
+
+    # 2. Train Model (CONSTRAINT MET: Train Once)
+    # We prepare training data using 'Preprocessing 1 + 2' 
+    train_data_path = prepare_training_data(file_pairs, preprocess_1_then_2, "prep12")
+    trained_model_path = train_model(train_data_path, "sam_finetuned_single")
+
+    # 3. Evaluate Trained Model (3 Iterations)
+    for prep_title, prep_func, prep_tag in combinations:
+        iou, f1 = run_inference_and_evaluate(
+            file_pairs, prep_func, prep_title, 
+            model_type="vit_b", model_checkpoint=trained_model_path, is_pretrained=False
+        )
+        results.append({"Pipeline": f"{prep_title} -> Trained Model", "F1": f1, "IoU@0.0": iou})
+
+    # 4. Output Final Table
+    print("\n### Final Evaluation Metrics (10 Image Subset)")
+    print("| Pipeline Combination | F1 Score | IoU@0.0 |")
+    print("| :--- | :--- | :--- |")
+    for res in results:
+        print(f"| {res['Pipeline']} | {res['F1']:.4f} | {res['IoU@0.0']:.4f} |")