Update app.py

app.py

@@ -1,74 +1,42 @@
-import os
-import subprocess
-
-# Clone the yolov5 repository and install its requirements
-if not os.path.exists('yolov5'):
-    subprocess.run(['git', 'clone', 'https://github.com/ultralytics/yolov5'], check=True)
-subprocess.run(['pip', 'install', '-r', 'yolov5/requirements.txt'], check=True)
-
 import torch
-import torchvision
-from torchvision.transforms import functional as F
-from PIL import Image
 import cv2
-import gradio as gr
 import numpy as np
-
-from yolov5.models.yolo import Model
+from PIL import Image
+from torchvision.transforms import functional as F
 from yolov5.utils.general import non_max_suppression
+from yolov5.models.yolo import Model
+import gradio as gr
 
+# Load YOLOv5 model
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-print(f"Using device: {device}")
 model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True).to(device)
 model.eval()
-print("Model loaded successfully")
 
 def preprocess_image(image):
-    try:
-        image = Image.fromarray(image)  # Convert numpy array to PIL Image
-        image_tensor = F.to_tensor(image).unsqueeze(0).to(device)
-        print(f"Preprocessed image tensor: {image_tensor.shape}")
-        return image_tensor
-    except Exception as e:
-        print(f"Error in preprocessing image: {e}")
-        return None
+    image = Image.fromarray(image)
+    image_tensor = F.to_tensor(image).unsqueeze(0).to(device)
+    return image_tensor
 
 def draw_boxes(image, outputs, threshold=0.3):
-    try:
-        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-        h, w, _ = image.shape
+    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    h, w, _ = image.shape
 
-        for box in outputs:
-            if box is not None:
-                x1, y1, x2, y2, score, label = box[:6]
-                if score > threshold:
-                    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
-                    cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
-                    text = f"{model.names[int(label)]:s}: {score:.2f}"
-                    cv2.putText(image, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
+    for box in outputs:
+        score, label, x1, y1, x2, y2 = box[4].item(), int(box[5].item()), box[0].item(), box[1].item(), box[2].item(), box[3].item()
+        if score > threshold:
+            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+            cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
+            text = f"{model.names[label]}: {score:.2f}"
+            cv2.putText(image, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
 
-        return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    except Exception as e:
-        print(f"Error in drawing boxes: {e}")
-        return image
+    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
 def detect_objects(image):
     image_tensor = preprocess_image(image)
-    if image_tensor is None:
-        return image
-    try:
-        outputs = model(image_tensor)
-        print(f"Model raw outputs: {outputs}")
-        outputs = non_max_suppression(outputs, conf_thres=0.25, iou_thres=0.45)
-        if outputs is None or len(outputs[0]) == 0:
-            print("No objects detected.")
-            return image
-        print(f"Filtered outputs: {outputs[0]}")
-        result_image = draw_boxes(image, outputs[0].cpu().numpy())
-        return result_image
-    except Exception as e:
-        print(f"Error in detecting objects: {e}")
-        return image
+    outputs = model(image_tensor)
+    outputs = non_max_suppression(outputs)[0]
+    result_image = draw_boxes(image, outputs.cpu().numpy())
+    return result_image
 
 iface = gr.Interface(
     fn=detect_objects,