Update app.py

app.py

@@ -1,91 +1,123 @@
-from segment_anything import sam_model_registry, SamPredictor
-import torch
 import cv2
 import numpy as np
-import gradio as gr
 import pandas as pd
+import torch
+import gradio as gr
 import matplotlib.pyplot as plt
+from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
+
+# Ensure the SAM model checkpoint is downloaded
+SAM_CHECKPOINT = "sam_vit_h.pth"  # Update path if needed
+MODEL_TYPE = "vit_h"
 
-# 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)
+# Check if CUDA is available for GPU processing
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load the SAM model
+try:
+    sam = sam_model_registry[MODEL_TYPE](checkpoint=SAM_CHECKPOINT).to(DEVICE)
+    mask_generator = SamAutomaticMaskGenerator(sam)
+except FileNotFoundError:
+    raise FileNotFoundError(f"Checkpoint file '{SAM_CHECKPOINT}' not found. Download it from: https://github.com/facebookresearch/segment-anything")
 
 def preprocess_image(image):
-    """Convert image to RGB format for SAM."""
+    """Convert image to grayscale and apply adaptive thresholding for better cell detection."""
     if len(image.shape) == 2:
         image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
-    return image
 
-def detect_blood_cells(image):
-    """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
+    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+
+    # Apply adaptive thresholding for better contrast
+    adaptive_thresh = cv2.adaptiveThreshold(
+        gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2
     )
 
-    contours_list = []
+    # Morphological operations to remove noise
+    kernel = np.ones((3, 3), np.uint8)
+    clean_mask = cv2.morphologyEx(adaptive_thresh, cv2.MORPH_CLOSE, kernel, iterations=2)
+    clean_mask = cv2.morphologyEx(clean_mask, cv2.MORPH_OPEN, kernel, iterations=2)
+
+    return clean_mask
+
+def detect_blood_cells(image):
+    """Detect blood cells using SAM segmentation + contour analysis."""
+    # Generate masks using SAM
+    masks = mask_generator.generate(image)
+
     features = []
+    processed_image = image.copy()
+
     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)
+        mask_binary = mask["segmentation"].astype(np.uint8) * 255  # Convert to binary
+        contours, _ = cv2.findContours(mask_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
 
-        for j, contour in enumerate(contours, 1):
+        for contour in contours:
             area = cv2.contourArea(contour)
             perimeter = cv2.arcLength(contour, True)
             circularity = 4 * np.pi * area / (perimeter * perimeter) if perimeter > 0 else 0
 
+            # Filter small or irregular shapes
             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
+                    features.append(
+                        {
+                            "label": len(features) + 1,
+                            "area": area,
+                            "perimeter": perimeter,
+                            "circularity": circularity,
+                            "centroid_x": cx,
+                            "centroid_y": cy,
+                        }
+                    )
+
+                    # Draw detected cell on image
+                    cv2.drawContours(processed_image, [contour], -1, (0, 255, 0), 2)
+                    cv2.putText(
+                        processed_image,
+                        str(len(features)),
+                        (cx, cy),
+                        cv2.FONT_HERSHEY_SIMPLEX,
+                        0.5,
+                        (0, 0, 255),
+                        1,
+                    )
+
+    return processed_image, features
 
 def process_image(image):
     if image is None:
         return None, None, None, None
 
-    contours, features, masks = detect_blood_cells(image)
-    vis_img = image.copy()
-
-    for feature in features:
-        contour = contours[int(feature['label'].split('-')[1]) - 1]
-        cv2.drawContours(vis_img, [contour], -1, (0, 255, 0), 2)
-        cv2.putText(vis_img, str(feature['label']), (feature['centroid_x'], feature['centroid_y']),
-                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
-
+    processed_img, features = detect_blood_cells(image)
     df = pd.DataFrame(features)
-    return vis_img, masks[0], df
+
+    return processed_img, df
 
 def analyze(image):
-    vis_img, mask, df = process_image(image)
+    processed_img, df = process_image(image)
 
-    plt.style.use('dark_background')
+    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].set_title('Cell Size Distribution')
+        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].set_title('Area vs Circularity')
+        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
+    return processed_img, fig, df
 
 # Gradio Interface
-demo = gr.Interface(fn=analyze, inputs=gr.Image(type="numpy"), outputs=[gr.Image(), gr.Image(), gr.Plot(), gr.Dataframe()])
+demo = gr.Interface(
+    fn=analyze,
+    inputs=gr.Image(type="numpy"),
+    outputs=[gr.Image(), gr.Plot(), gr.Dataframe()],
+    title="Blood Cell Detection",
+    description="Detect and analyze blood cells using SAM segmentation & contour analysis.",
+)
+
 demo.launch()