SAHI

app.py

@@ -1,69 +1,108 @@
 import os
-# os.system('')
 import cv2
 import gradio as gr
 import torch
 from ultralytics import YOLO
-import random
-import numpy as np
-random.seed(42)
-np.random.seed(42)
+from sahi import AutoDetectionModel
+from sahi.predict import get_sliced_prediction
+import time
 
+# Configuration
+MODEL_PATH = os.getenv("MODEL_PATH", "last.pt")
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+VALID_EXTENSIONS = [".jpg", ".jpeg", ".png"]
+
+# Load models
 try:
-    model = YOLO("last.pt")
+    yolo_model = YOLO(MODEL_PATH).to(DEVICE)
     print("YOLO model loaded successfully.")
-except FileNotFoundError:
-    print("Error: 'yolo_modeln11_1502.pt' not found.")
-    model = None
 except Exception as e:
-    print(f"An error occurred while loading the model: {e}")
-    model = None
-
-# Function to predict and show bounding boxes
-def predict_and_show_bounding_boxes(image_path):
-    if model is None:
-        return None, "Error: Model not loaded."
-    
-    try:
-        # Load the image using cv2
-        img = cv2.imread(image_path)
-        if img is None:
-            print(f"Error: Could not load image at {image_path}")
-            return None, "Error: Could not load image"
+    print(f"Error loading YOLO model: {e}")
+    yolo_model = None
 
-        # Perform inference using the YOLO model
-        results = model(image_path,conf=0.5)[0]
-        boxes = results.boxes
+try:
+    sahi_model = AutoDetectionModel.from_pretrained(
+        model_type="ultralytics",
+        model_path=MODEL_PATH,
+        confidence_threshold=0.5,
+        device=DEVICE,
+    )
+    print("SAHI model loaded successfully.")
+except Exception as e:
+    print(f"Error loading SAHI model: {e}")
+    sahi_model = None
 
-        if len(boxes) == 0:
-            # No defects found, show the zero defects image
-            zero_defects_img = cv2.imread('zero_defects.png')
-            if zero_defects_img is not None:
-                return zero_defects_img
-            else:
-                return None, "Error: Could not load zero defects image"
+def predict_and_show_bounding_boxes(image_path, model_choice, conf_threshold=0.5):
+    if not image_path or not any(image_path.lower().endswith(ext) for ext in VALID_EXTENSIONS):
+        return None, "Error: Invalid or unsupported image format."
 
-        for box in boxes:
-            xyxy = box.xyxy[0].tolist()
-            x_min, y_min, x_max, y_max = int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])
-            conf = box.conf[0].item()
-            cls = int(box.cls[0])
-            cv2.rectangle(img, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
-            label = f"{results.names[cls]}: {conf:.2f}"
-            cv2.putText(img, label, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+    # Resize image for faster processing
+    img = cv2.imread(image_path)
+    if img is None:
+        return None, "Error: Could not load image."
+    img = cv2.resize(img, (640, 640))  # Resize to 640x640
 
-        return img
-    except Exception as e:
-        print(f"An error occurred during prediction: {e}")
-        return None, str(e)
+    if model_choice == "YOLO":
+        if yolo_model is None:
+            return None, "Error: YOLO model not loaded."
+        try:
+            start_time = time.time()
+            results = yolo_model(img, conf=conf_threshold)[0]
+            print(f"YOLO inference time: {time.time() - start_time:.2f} seconds")
+            boxes = results.boxes
+            if len(boxes) == 0:
+                return cv2.cvtColor(img, cv2.COLOR_BGR2RGB), "No defects detected."
+            for box in boxes:
+                xyxy = box.xyxy[0].tolist()
+                x_min, y_min, x_max, y_max = map(int, xyxy[:4])
+                conf = box.conf[0].item()
+                cls = int(box.cls[0])
+                cv2.rectangle(img, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
+                label = f"{results.names[cls]}: {conf:.2f}"
+                cv2.putText(img, label, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+            return cv2.cvtColor(img, cv2.COLOR_BGR2RGB), "Detection complete."
+        except Exception as e:
+            return None, f"Error during YOLO prediction: {e}"
+    elif model_choice == "SAHI":
+        if sahi_model is None:
+            return None, "Error: SAHI model not loaded."
+        try:
+            start_time = time.time()
+            result = get_sliced_prediction(
+                img,  # Use resized image
+                sahi_model,
+                slice_height=512,
+                slice_width=512,
+                overlap_height_ratio=0.1,
+                overlap_width_ratio=0.1,
+            )
+            print(f"SAHI inference time: {time.time() - start_time:.2f} seconds")
+            if len(result.object_prediction_list) == 0:
+                return cv2.cvtColor(img, cv2.COLOR_BGR2RGB), "No defects detected."
+            for pred in result.object_prediction_list:
+                box = pred.bbox.to_xyxy()
+                x_min, y_min, x_max, y_max = map(int, box)
+                label = f"{pred.category.name}: {pred.score.value:.2f}"
+                cv2.rectangle(img, (x_min, y_min), (x_max, y_max), (255, 0, 0), 2)
+                cv2.putText(img, label, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 0, 0), 2)
+            return cv2.cvtColor(img, cv2.COLOR_BGR2RGB), "Detection complete."
+        except Exception as e:
+            return None, f"Error during SAHI prediction: {e}"
+    return None, "Invalid model choice."
 
-# Create Gradio interface
+# Gradio interface
 iface = gr.Interface(
     fn=predict_and_show_bounding_boxes,
-    inputs=gr.Image(type="filepath"),
-    outputs=[gr.Image()],
-    title="Defect Detection",
-    description="Upload an image to detect defects"
+    inputs=[
+        gr.Image(type="filepath", label="Upload Image"),
+        gr.Radio(choices=["YOLO", "SAHI"], label="Choose Detection Mode", value="YOLO"),
+        gr.Slider(minimum=0.1, maximum=0.9, value=0.5, label="Confidence Threshold"),
+    ],
+    outputs=[gr.Image(label="Result", width=640, height=640), gr.Textbox(label="Message")],
+    title="PCB Defect Detection",
+    description="Upload a PCB image and choose YOLO (green boxes) or SAHI (red boxes) for defect detection. Adjust confidence threshold for sensitivity.",
 )
 
-iface.launch(share=True)
+if __name__ == "__main__":
+    share = os.getenv("HF_SHARE", "False").lower() == "true"
+    iface.launch(share=share)