Update app.py

app.py

@@ -1,9 +1,16 @@
-import gradio as gr
+import os
 import torch
+import gradio as gr
 import cv2
 import numpy as np
 from PIL import Image
 from torchvision.transforms import functional as F
+
+# Clone yolov5 if not present
+if not os.path.exists("yolov5"):
+    os.system("git clone https://github.com/ultralytics/yolov5.git")
+    os.system("pip install -r yolov5/requirements.txt")
+
 from yolov5.utils.general import non_max_suppression
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -17,7 +24,6 @@ def preprocess_image(image):
 
 def draw_boxes(image, outputs, threshold=0.3):
     image = np.array(image)
-    h, w, _ = image.shape
     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:
@@ -33,13 +39,8 @@ def detect_objects(image):
     outputs = non_max_suppression(outputs)[0]
     return draw_boxes(image, outputs)
 
-iface = gr.Interface(
-    fn=detect_objects,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="YOLO Object Detector",
-    description="Upload an image to detect objects using YOLOv5."
-)
+iface = gr.Interface(fn=detect_objects, inputs=gr.Image(type="pil"), outputs=gr.Image(type="pil"),
+                     title="YOLO Object Detector", description="Upload an image to detect objects using YOLOv5.")
 
 if __name__ == "__main__":
     iface.launch()