Update app.py

app.py

@@ -1,117 +1,70 @@
 import os
+os.system('pip install --upgrade gradio')
+# os.system('')
 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)
 
-# Configuration
-MODEL_PATH = os.getenv("MODEL_PATH", "last.pt")
-ZERO_DEFECTS_PATH = "zero_defects.png"
-VALID_EXTENSIONS = [".jpg", ".jpeg", ".png"]
-
-# Load YOLO model
 try:
-    if not os.path.exists(MODEL_PATH):
-        raise FileNotFoundError(f"Model file {MODEL_PATH} not found.")
-    yolo_model = YOLO(MODEL_PATH)
+    model = YOLO("last.pt")
     print("YOLO model loaded successfully.")
+except FileNotFoundError:
+    print("Error: 'yolo_modeln11_1502.pt' not found.")
+    model = None
 except Exception as e:
-    print(f"Error loading YOLO model: {e}")
-    yolo_model = None
+    print(f"An error occurred while loading the model: {e}")
+    model = None
 
-# Load SAHI model
-try:
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    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
+# 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"
 
-def predict_and_show_bounding_boxes(image_path, model_choice, conf_threshold=0.5):
-    # Validate image path
-    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)
 
-    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=conf_threshold)[0]
-            boxes = results.boxes
-            if len(boxes) == 0:
-                if os.path.exists(ZERO_DEFECTS_PATH):
-                    return cv2.imread(ZERO_DEFECTS_PATH), "No defects detected."
-                return img, "No defects detected (zero_defects.png not found)."
-            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:
-            img = cv2.imread(image_path)
-            if img is None:
-                return None, "Error: Could not load image."
-            result = get_sliced_prediction(
-                image_path,
-                sahi_model,
-                slice_height=256,
-                slice_width=256,
-                overlap_height_ratio=0.2,
-                overlap_width_ratio=0.2,
-            )
-            if len(result.object_prediction_list) == 0:
-                if os.path.exists(ZERO_DEFECTS_PATH):
-                    return cv2.imread(ZERO_DEFECTS_PATH), "No defects detected."
-                return img, "No defects detected (zero_defects.png not found)."
-            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."
+        return img
+    except Exception as e:
+        print(f"An error occurred during prediction: {e}")
+        return None, str(e)
 
-# Gradio interface
+# Create Gradio interface
 iface = gr.Interface(
     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"),
-        gr.Slider(minimum=0.1, maximum=0.9, value=0.5, label="Confidence Threshold"),
-    ],
-    outputs=[gr.Image(label="Result"), 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.",
+    inputs=gr.Image(type="filepath"),
+    outputs=[gr.Image()],
+    title="Defect Detection",
+    description="Upload an image to detect defects"
 )
 
-if __name__ == "__main__":
-    share = os.getenv("HF_SHARE", "False").lower() == "true"
-    iface.launch(share=share)
+iface.launch(share=True)