Update app.py

app.py

@@ -6,24 +6,25 @@ from PIL import Image
 import random
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model = torch.hub.load('ultralytics/yolov5', 'yolov5x', pretrained=True).to(device)  # Load YOLOv5x model
 
-CLASS_NAMES = model.names  
+# Use a smaller model for faster inference
+model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True).to(device)
+model.eval()
 
+CLASS_NAMES = model.names  
 random.seed(42)
 CLASS_COLORS = {cls: (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) for cls in CLASS_NAMES}
 
 def preprocess_image(image):
-    image = Image.fromarray(image)
-    image = image.convert("RGB")
+    image = Image.fromarray(image).convert("RGB").resize((640, 640))
     return image
 
 def detect_objects(image):
     image = preprocess_image(image)
 
-    results = model(image)
+    results = model([image])  # Batch processing for efficiency
 
-    image = np.array(image)
+    image = np.array(image)  
     
     for *box, conf, cls in results.xyxy[0]:  
         x1, y1, x2, y2 = map(int, box)
@@ -33,15 +34,12 @@ def detect_objects(image):
         color = CLASS_COLORS[class_name]
 
         cv2.rectangle(image, (x1, y1), (x2, y2), color, 4)  
-
-        
         label = f"{class_name} ({confidence:.1f}%)"
         cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 
-                    1, color, 3, cv2.LINE_AA)  # Larger text
+                    1, color, 3, cv2.LINE_AA)
 
     return image
 
-
 iface = gr.Interface(
     fn=detect_objects,
     inputs=gr.Image(type="numpy", label="Upload Image"),
@@ -49,7 +47,7 @@ iface = gr.Interface(
     title="Object Detection with YOLOv5",
     description="Use webcam or upload an image to detect objects.",
     allow_flagging="never",
-    examples=["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
+    examples=["examples/spring_street_after.jpg", "examples/pexels-hikaique-109919.jpg"]
 )
 
 iface.launch()