Update app.py

app.py

@@ -1,123 +1,45 @@
+import gradio as gr
 import cv2
 import numpy as np
-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"
-
-# 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 grayscale and apply adaptive thresholding for better cell detection."""
-    if len(image.shape) == 2:
-        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
-
-    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
-    )
-
-    # 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_binary = mask["segmentation"].astype(np.uint8) * 255  # Convert to binary
-        contours, _ = cv2.findContours(mask_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-
-        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": 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
-
-    processed_img, features = detect_blood_cells(image)
-    df = pd.DataFrame(features)
-
-    return processed_img, df
-
-def analyze(image):
-    processed_img, 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].set_title("Cell Size Distribution")
-
-        axes[1].scatter(df["area"], df["circularity"], alpha=0.6, c="magenta")
-        axes[1].set_title("Area vs Circularity")
-
-    return processed_img, fig, df
+from collections import Counter
+from ultralytics import YOLO
+from huggingface_hub import hf_hub_download
+
+# Download model from Hugging Face repo
+MODEL_PATH = hf_hub_download(
+    repo_id="ibrahim313/Bioengineering_Query_Tool_image_based",
+    filename="best.pt"
+)
 
-# Gradio Interface
-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.",
+# Load the YOLOv10 model
+model = YOLO(MODEL_PATH)
+
+def predict(image):
+    # Convert the image from BGR to RGB
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    
+    # Perform prediction
+    results = model.predict(source=image_rgb, imgsz=640, conf=0.25)
+    
+    # Get the annotated image
+    annotated_img = results[0].plot()
+    
+    # Extract detection data
+    detections = results[0].boxes.data if results[0].boxes is not None else []
+    class_names = [model.names[int(cls)] for cls in detections[:, 5]] if len(detections) > 0 else []
+    count = Counter(class_names)
+
+    # Create a string representation of the detections
+    detection_str = ', '.join([f"{name}: {count}" for name, count in count.items()]) if class_names else "No detections"
+
+    return annotated_img, detection_str 
+
+app = gr.Interface(
+    predict, 
+    inputs=gr.Image(type="numpy", label="Upload an Image"),
+    outputs=[gr.Image(type="numpy", label="Annotated Image"), gr.Textbox(label="Detection Counts")],
+    title="Blood Cell Count", 
+    description="Upload an image and YOLOv10 will detect blood cells."
 )
 
-demo.launch()
+app.launch()