Update app.py

app.py

@@ -1,109 +1,91 @@
+from segment_anything import sam_model_registry, SamPredictor
+import torch
 import cv2
 import numpy as np
-import pandas as pd
 import gradio as gr
+import pandas as pd
 import matplotlib.pyplot as plt
-from datetime import datetime
+
+# Load SAM model
+sam_checkpoint = "sam_vit_h.pth"  # Checkpoint file (download it from Meta AI)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model_type = "vit_h"
+sam = sam_model_registry[model_type](checkpoint=sam_checkpoint).to(device)
+predictor = SamPredictor(sam)
 
 def preprocess_image(image):
-    """Enhance image contrast, apply thresholding, and clean noise."""
+    """Convert image to RGB format for SAM."""
     if len(image.shape) == 2:
         image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
-
-    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
-    
-    # Apply Gaussian blur to remove noise
-    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
-    
-    # Otsu's Thresholding (more robust than adaptive for blood cells)
-    _, binary = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
-
-    # Morphological operations to improve segmentation
-    kernel = np.ones((3, 3), np.uint8)
-    clean_mask = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel, iterations=2)
-    clean_mask = cv2.morphologyEx(clean_mask, cv2.MORPH_OPEN, kernel, iterations=1)
-
-    return clean_mask
+    return image
 
 def detect_blood_cells(image):
-    """Detect blood cells and extract features."""
-    mask = preprocess_image(image)
-    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-
+    """Detect blood cells using SAM."""
+    image = preprocess_image(image)
+    predictor.set_image(image)
+
+    # Generate automatic masks (SAM can also take prompts for guided segmentation)
+    masks, _, _ = predictor.predict(
+        point_coords=None, 
+        point_labels=None, 
+        multimask_output=True
+    )
+
+    contours_list = []
     features = []
-    total_area = 0
-
-    for i, contour in enumerate(contours, 1):
-        area = cv2.contourArea(contour)
-        perimeter = cv2.arcLength(contour, True)
-        circularity = (4 * np.pi * area / (perimeter * perimeter)) if perimeter > 0 else 0
-
-        # Filtering: Only count reasonable-sized circular objects
-        if 100 < area < 5000 and circularity > 0.7:
-            M = cv2.moments(contour)
-            if M["m00"] != 0:
-                cx = int(M["m10"] / M["m00"])
-                cy = int(M["m01"] / M["m00"])
-                features.append({
-                    'ID': i, 'Area': area, 'Perimeter': perimeter, 
-                    'Circularity': circularity, 'Centroid_X': cx, 'Centroid_Y': cy
-                })
-                total_area += area
-
-    # Summary Statistics
-    avg_cell_size = total_area / len(features) if features else 0
-    cell_density = len(features) / (image.shape[0] * image.shape[1])  # Density per pixel
-
-    summary = {
-        'Total Cells': len(features),
-        'Avg Cell Size': avg_cell_size,
-        'Cell Density': cell_density
-    }
-
-    return contours, features, mask, summary
+    for i, mask in enumerate(masks):
+        mask = mask.astype(np.uint8) * 255  # Convert boolean mask to uint8
+        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+        for j, contour in enumerate(contours, 1):
+            area = cv2.contourArea(contour)
+            perimeter = cv2.arcLength(contour, True)
+            circularity = 4 * np.pi * area / (perimeter * perimeter) if perimeter > 0 else 0
+
+            if 100 < area < 5000 and circularity > 0.7:
+                M = cv2.moments(contour)
+                if M["m00"] != 0:
+                    cx = int(M["m10"] / M["m00"])
+                    cy = int(M["m01"] / M["m00"])
+                    features.append({
+                        'label': f"{i}-{j}", 'area': area, 'perimeter': perimeter,
+                        'circularity': circularity, 'centroid_x': cx, 'centroid_y': cy
+                    })
+                contours_list.append(contour)
+
+    return contours_list, features, masks
 
 def process_image(image):
     if image is None:
-        return None, None, None, None, None
+        return None, None, None, None
 
-    contours, features, mask, summary = detect_blood_cells(image)
+    contours, features, masks = detect_blood_cells(image)
     vis_img = image.copy()
 
     for feature in features:
-        contour = contours[feature['ID'] - 1]
+        contour = contours[int(feature['label'].split('-')[1]) - 1]
         cv2.drawContours(vis_img, [contour], -1, (0, 255, 0), 2)
-        cv2.putText(vis_img, str(feature['ID']), (feature['Centroid_X'], feature['Centroid_Y']),
+        cv2.putText(vis_img, str(feature['label']), (feature['centroid_x'], feature['centroid_y']),
                     cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
 
     df = pd.DataFrame(features)
-    return vis_img, mask, df, summary
+    return vis_img, masks[0], df
 
 def analyze(image):
-    vis_img, mask, df, summary = process_image(image)
+    vis_img, mask, df = process_image(image)
 
     plt.style.use('dark_background')
     fig, axes = plt.subplots(1, 2, figsize=(12, 5))
 
     if not df.empty:
-        axes[0].hist(df['Area'], bins=20, color='cyan', edgecolor='black')
+        axes[0].hist(df['area'], bins=20, color='cyan', edgecolor='black')
         axes[0].set_title('Cell Size Distribution')
 
-        axes[1].scatter(df['Area'], df['Circularity'], alpha=0.6, c='magenta')
+        axes[1].scatter(df['area'], df['circularity'], alpha=0.6, c='magenta')
         axes[1].set_title('Area vs Circularity')
 
-    return vis_img, mask, fig, df, summary
+    return vis_img, mask, fig, df
 
 # Gradio Interface
-demo = gr.Interface(
-    fn=analyze,
-    inputs=gr.Image(type="numpy"),
-    outputs=[
-        gr.Image(label="Processed Image"),
-        gr.Image(label="Binary Mask"),
-        gr.Plot(label="Analysis Plots"),
-        gr.Dataframe(label="Detected Cells Data"),
-        gr.JSON(label="Summary Statistics")
-    ]
-)
-
+demo = gr.Interface(fn=analyze, inputs=gr.Image(type="numpy"), outputs=[gr.Image(), gr.Image(), gr.Plot(), gr.Dataframe()])
 demo.launch()