Add ground truth comparison: compare SAM 3 segmentation with ground truth masks (BraTS, Kaggle datasets)

app.py

@@ -93,8 +93,73 @@ except:
     except Exception as e:
         print(f"⚠️ Could not create demo file: {e}")
 
-def process_medical_image(image_file, prompt_text, modality, window_type):
-    """Process a DICOM or standard image file (PNG/JPG) and perform segmentation using SAM 3."""
+def compare_with_ground_truth(pred_mask, gt_mask_path):
+    """Compare SAM 3 prediction with ground truth mask and return comparison metrics."""
+    try:
+        gt_mask = Image.open(gt_mask_path)
+        gt_array = np.array(gt_mask.convert('L')) > 127  # Binarize
+        
+        # Resize prediction mask to match ground truth if needed
+        if pred_mask.shape != gt_array.shape:
+            from PIL import Image as PILImage
+            pred_pil = PILImage.fromarray((pred_mask * 255).astype(np.uint8))
+            pred_pil = pred_pil.resize(gt_mask.size, PILImage.NEAREST)
+            pred_mask = np.array(pred_pil) > 127
+        
+        # Calculate metrics
+        intersection = np.logical_and(pred_mask, gt_array).sum()
+        union = np.logical_or(pred_mask, gt_array).sum()
+        dice_score = (2.0 * intersection) / (pred_mask.sum() + gt_array.sum()) if (pred_mask.sum() + gt_array.sum()) > 0 else 0.0
+        iou_score = intersection / union if union > 0 else 0.0
+        
+        # Create comparison visualization
+        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+        
+        axes[0].imshow(pred_mask, cmap='spring')
+        axes[0].set_title('SAM 3 Prediction')
+        axes[0].axis('off')
+        
+        axes[1].imshow(gt_array, cmap='cool')
+        axes[1].set_title('Ground Truth')
+        axes[1].axis('off')
+        
+        # Overlay comparison
+        comparison = np.zeros((*pred_mask.shape, 3))
+        comparison[pred_mask & gt_array] = [0, 1, 0]  # Green: True Positive
+        comparison[pred_mask & ~gt_array] = [1, 0, 0]  # Red: False Positive
+        comparison[~pred_mask & gt_array] = [0, 0, 1]  # Blue: False Negative
+        
+        axes[2].imshow(comparison)
+        axes[2].set_title(f'Comparison\nDice: {dice_score:.3f}, IoU: {iou_score:.3f}')
+        axes[2].axis('off')
+        
+        plt.tight_layout()
+        
+        output_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
+        output_path = output_file.name
+        output_file.close()
+        
+        plt.savefig(output_path, bbox_inches='tight', dpi=100)
+        plt.close()
+        
+        return output_path, dice_score, iou_score
+    except Exception as e:
+        print(f"⚠️ Error comparing with ground truth: {e}")
+        return None, 0.0, 0.0
+
+def process_medical_image(image_file, prompt_text, modality, window_type, return_mask=False):
+    """Process a DICOM or standard image file (PNG/JPG) and perform segmentation using SAM 3.
+    
+    Args:
+        image_file: Path to image file
+        prompt_text: Text prompt for segmentation
+        modality: CT or MRI
+        window_type: Windowing strategy
+        return_mask: If True, also return the binary mask array
+    
+    Returns:
+        Path to output image, and optionally the mask array
+    """
     if model is None or processor is None:
         print("❌ Error: Model not loaded.")
         return None
@@ -223,6 +288,7 @@ def process_medical_image(image_file, prompt_text, modality, window_type):
         plt.figure(figsize=(10, 10))
         plt.imshow(pil_image)
         
+        final_mask = None
         if 'masks' in results and results['masks'] is not None:
             masks = results['masks'].cpu().numpy()
             if len(masks) > 0:
@@ -243,6 +309,8 @@ def process_medical_image(image_file, prompt_text, modality, window_type):
         plt.savefig(output_path, bbox_inches='tight', pad_inches=0, dpi=100)
         plt.close()
         
+        if return_mask:
+            return output_path, final_mask
         return output_path
         
     except pydicom.errors.InvalidDicomError as e:
@@ -279,6 +347,32 @@ def process_with_status(image_file, prompt_text, modality, window_type):
     else:
         return result, "✅ Segmentation complete!"
 
+def process_with_ground_truth(image_file, gt_mask_file, prompt_text, modality, window_type):
+    """Process image and compare with ground truth segmentation mask."""
+    if model is None or processor is None:
+        return None, None, 0.0, 0.0, "❌ Error: Model not loaded."
+    
+    if image_file is None:
+        return None, None, 0.0, 0.0, "⚠️ Please upload a medical image file."
+    
+    if gt_mask_file is None:
+        return None, None, 0.0, 0.0, "⚠️ Please upload a ground truth mask file."
+    
+    # Process image and get mask
+    result, pred_mask = process_medical_image(image_file, prompt_text, modality, window_type, return_mask=True)
+    
+    if result is None or pred_mask is None:
+        return None, None, 0.0, 0.0, "❌ Processing failed. Check console for error details."
+    
+    # Compare with ground truth
+    comparison_path, dice_score, iou_score = compare_with_ground_truth(pred_mask, gt_mask_file)
+    
+    if comparison_path:
+        status = f"✅ Segmentation complete!\nDice Score: {dice_score:.3f}\nIoU Score: {iou_score:.3f}"
+        return result, comparison_path, dice_score, iou_score, status
+    else:
+        return result, None, 0.0, 0.0, "✅ Segmentation complete, but comparison failed."
+
 def process_sequence(image_files, prompt_text, modality, window_type):
     """Process multiple images from the same subject and return gallery of results."""
     if model is None or processor is None:
@@ -445,6 +539,78 @@ with gr.Blocks() as demo:
                         interactive=False,
                         lines=5
                     )
+        
+        with gr.Tab("Compare with Ground Truth"):
+            gr.Markdown("**Compare SAM 3 segmentation with ground truth masks (e.g., from BraTS, Kaggle datasets)**")
+            with gr.Row():
+                with gr.Column():
+                    file_input_gt = gr.File(
+                        label="Upload Medical Image (DICOM .dcm, PNG, JPG)", 
+                        file_types=[".dcm", ".png", ".jpg", ".jpeg"],
+                        type="filepath"
+                    )
+                    
+                    gt_mask_input = gr.File(
+                        label="Upload Ground Truth Mask (PNG, JPG)", 
+                        file_types=[".png", ".jpg", ".jpeg"],
+                        type="filepath",
+                        info="Upload the ground truth segmentation mask (binary or labeled image)"
+                    )
+                    
+                    text_input_gt = gr.Textbox(
+                        label="Text Prompt", 
+                        value="brain", 
+                        placeholder="e.g. brain, tumor, skull",
+                        info="Describe what anatomical structure or region you want to segment"
+                    )
+                    
+                    with gr.Row():
+                        modality_dropdown_gt = gr.Dropdown(
+                            ["CT", "MRI"], 
+                            label="Modality", 
+                            value="MRI",
+                            info="Select the imaging modality"
+                        )
+                        window_dropdown_gt = gr.Dropdown(
+                            ["Brain (Grey Matter)", "Bone (Skull)", "Soft Tissue (Face)"], 
+                            label="Windowing Strategy (CT only)", 
+                            value="Brain (Grey Matter)",
+                            info="CT windowing preset (ignored for MRI)"
+                        )
+                    
+                    submit_gt_btn = gr.Button("Compare Segmentation", variant="primary", size="lg")
+                    
+                with gr.Column():
+                    image_output_gt = gr.Image(
+                        label="SAM 3 Segmentation",
+                        type="filepath"
+                    )
+                    
+                    comparison_output = gr.Image(
+                        label="Comparison: SAM 3 vs Ground Truth",
+                        type="filepath"
+                    )
+                    
+                    gr.Markdown("### Metrics")
+                    dice_score_text = gr.Textbox(
+                        label="Dice Score",
+                        value="--",
+                        interactive=False
+                    )
+                    
+                    iou_score_text = gr.Textbox(
+                        label="IoU Score",
+                        value="--",
+                        interactive=False
+                    )
+                    
+                    gr.Markdown("### Status")
+                    status_gt_text = gr.Textbox(
+                        label="Processing Status",
+                        value="Ready. Upload image and ground truth mask to compare.",
+                        interactive=False,
+                        lines=3
+                    )
 
     # Single image processing
     load_demo_btn.click(
@@ -465,6 +631,13 @@ with gr.Blocks() as demo:
         inputs=[files_input, text_input_batch, modality_dropdown_batch, window_dropdown_batch],
         outputs=[gallery_output, status_batch_text]
     )
+    
+    # Ground truth comparison
+    submit_gt_btn.click(
+        fn=process_with_ground_truth,
+        inputs=[file_input_gt, gt_mask_input, text_input_gt, modality_dropdown_gt, window_dropdown_gt],
+        outputs=[image_output_gt, comparison_output, dice_score_text, iou_score_text, status_gt_text]
+    )
 
 if __name__ == "__main__":
     demo.launch()