Update app.py

app.py

@@ -1,55 +1,119 @@
+import os
+os.system('pip install --upgrade gradio sahi')
+import cv2
 import gradio as gr
-from sahi.predict import get_sliced_prediction
-from sahi.model import Yolov5DetectionModel
-from PIL import Image
-import numpy as np
 import torch
 from ultralytics import YOLO
+from sahi import AutoDetectionModel
+from sahi.predict import get_sliced_prediction
+import random
+import numpy as np
 
-# Initialize YOLOv11n model
-yolo_model = YOLO('yolov11n.pt')
-print("YOLO model loaded successfully.")
+random.seed(42)
+np.random.seed(42)
 
-# Wrap with SAHI
-sahi_model = Yolov5DetectionModel(
-    model_path='yolov11n.pt',
-    confidence_threshold=0.3,
-    device='cuda' if torch.cuda.is_available() else 'cpu',
-)
-print("SAHI model loaded successfully.")
-
-# Inference function
-def detect_objects(img: Image.Image):
-    image_np = np.array(img)
-    result = get_sliced_prediction(
-        image_np,
-        detection_model=sahi_model,
-        slice_height=512,
-        slice_width=512,
-        overlap_height_ratio=0.2,
-        overlap_width_ratio=0.2
+# Load default YOLO model
+try:
+    yolo_model = YOLO("last.pt")
+    print("YOLO model loaded successfully.")
+except FileNotFoundError:
+    print("Error: 'last.pt' not found.")
+    yolo_model = None
+except Exception as e:
+    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)
 
-    detections = result.object_prediction_list
+            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")
 
-    # Draw boxes
-    from PIL import ImageDraw
-    draw = ImageDraw.Draw(img)
+            return img
+        except Exception as e:
+            return None, str(e)
 
-    for det in detections:
-        box = det.bbox.to_xyxy()
-        label = det.category.name
-        draw.rectangle(box, outline="red", width=2)
-        draw.text((box[0], box[1]), label, fill="red")
+    else:
+        return None, "Invalid model choice."
 
-    return img
 
-# Create Gradio interface
+# Gradio interface
 iface = gr.Interface(
-    fn=detect_objects,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="YOLOv11n SAHI Object Detection",
+    fn=predict_and_show_bounding_boxes,
+    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 and choose YOLO or SAHI for defect detection.",
 )
 
-iface.launch()
+iface.launch(share=True)