Update app.py

app.py

@@ -26,7 +26,7 @@ else:
 # Optimize model for speed
 model.conf = 0.3  # Lower confidence threshold
 model.iou = 0.3   # Non-Maximum Suppression IoU threshold
-model.classes = None  # Detect all classes
+model.classes = None  # Detect all 80+ COCO classes
 
 if device.type == "cuda":
     model.half()  # Use FP16 for faster inference
@@ -37,7 +37,7 @@ model.eval()
 
 # Pre-generate colors for bounding boxes
 np.random.seed(42)
-colors = np.random.uniform(0, 255, size=(len(model.names), 3))
+colors = np.random.randint(0, 255, size=(len(model.names), 3), dtype=np.uint8)
 
 def process_video(video_path):
     cap = cv2.VideoCapture(video_path)
@@ -121,8 +121,13 @@ with gr.Blocks(title="Real-Time YOLOv5 Video & Image Detection") as demo:
             process_button.click(fn=process_video, inputs=video_input, outputs=video_output)
 
         with gr.TabItem("Image Detection"):
-            image_input = gr.Image(type="pil", label="Upload Image")
-            image_output = gr.Image(label="Detected Objects")
-            image_input.change(fn=process_image, inputs=image_input, outputs=image_output)
+            with gr.Row():
+                image_input = gr.Image(type="pil", label="Upload Image", width=256, height=256)
+            with gr.Row():
+                submit_button = gr.Button("Submit")
+                clear_button = gr.Button("Clear")
+            image_output = gr.Image(label="Detected Objects", width=256, height=256)
+            submit_button.click(fn=process_image, inputs=image_input, outputs=image_output)
+            clear_button.click(fn=lambda: None, inputs=None, outputs=image_output)
 
 demo.launch()