Update app.py

app.py

@@ -1,70 +1,119 @@
 import os
-os.system('pip install --upgrade gradio')
-# os.system('')
+os.system('pip install --upgrade gradio sahi')
 import cv2
 import gradio as gr
 import torch
 from ultralytics import YOLO
+from sahi import AutoDetectionModel
+from sahi.predict import get_sliced_prediction
 import random
 import numpy as np
+
 random.seed(42)
 np.random.seed(42)
 
+# Load default YOLO model
 try:
-    model = YOLO("last.pt")
+    yolo_model = YOLO("last.pt")
     print("YOLO model loaded successfully.")
 except FileNotFoundError:
-    print("Error: 'yolo_modeln11_1502.pt' not found.")
-    model = None
+    print("Error: 'last.pt' not found.")
+    yolo_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"
-
-        # Perform inference using the YOLO model
-        results = model(image_path,conf=0.5)[0]
-        boxes = results.boxes
-
-        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"
-
-        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)
-
-        return img
-    except Exception as e:
-        print(f"An error occurred during prediction: {e}")
-        return None, str(e)
-
-# Create Gradio interface
+    print(f"An error occurred while loading the YOLO model: {e}")
+    yolo_model = None
+
+# Load SAHI model
+try:
+    sahi_model = AutoDetectionModel.from_pretrained(
+        model_type="ultralytics",
+        model_path="last.pt",  # same model used for consistency
+        confidence_threshold=0.5,
+        device="cpu",
+    )
+    print("SAHI model loaded successfully.")
+except Exception as e:
+    print(f"An error occurred while loading the SAHI model: {e}")
+    sahi_model = None
+
+# Prediction function
+def predict_and_show_bounding_boxes(image_path, model_choice):
+    if model_choice == "YOLO":
+        if yolo_model is None:
+            return None, "Error: YOLO model not loaded."
+
+        try:
+            img = cv2.imread(image_path)
+            if img is None:
+                return None, "Error: Could not load image"
+
+            results = yolo_model(image_path, conf=0.5)[0]
+            boxes = results.boxes
+
+            if len(boxes) == 0:
+                zero_defects_img = cv2.imread('zero_defects.png')
+                return zero_defects_img if zero_defects_img is not None else (None, "Error: Could not load zero defects image")
+
+            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 img
+        except Exception as e:
+            return None, str(e)
+
+    elif model_choice == "SAHI":
+        if sahi_model is None:
+            return None, "Error: SAHI model not loaded."
+
+        try:
+            result = get_sliced_prediction(
+                image_path,
+                sahi_model,
+                slice_height=256,
+                slice_width=256,
+                overlap_height_ratio=0.2,
+                overlap_width_ratio=0.2,
+            )
+
+            img = cv2.imread(image_path)
+            if img is None:
+                return None, "Error: Could not load image"
+
+            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)
+
+            if len(result.object_prediction_list) == 0:
+                zero_defects_img = cv2.imread('zero_defects.png')
+                return zero_defects_img if zero_defects_img is not None else (None, "Error: Could not load zero defects image")
+
+            return img
+        except Exception as e:
+            return None, str(e)
+
+    else:
+        return None, "Invalid model choice."
+
+
+# Gradio interface
 iface = gr.Interface(
     fn=predict_and_show_bounding_boxes,
-    inputs=gr.Image(type="filepath"),
-    outputs=[gr.Image()],
+    inputs=[
+        gr.Image(type="filepath", label="Upload Image"),
+        gr.Radio(choices=["YOLO", "SAHI"], label="Choose Detection Mode", value="YOLO"),
+    ],
+    outputs=[gr.Image(label="Result"), gr.Textbox(label="Message")],
     title="Defect Detection",
-    description="Upload an image to detect defects"
+    description="Upload an image and choose YOLO or SAHI for defect detection.",
 )
 
-iface.launch(share=True)
+iface.launch(share=True)