Update app.py

app.py

@@ -5,14 +5,14 @@ import numpy as np
 import gradio as gr
 from PIL import Image
 
-# Load YOLOv5 model
+# Load YOLOv8 model
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model = YOLO("yolov5s.pt")  # Load pre-trained YOLOv5s model
+model = YOLO("yolov8x.pt")  # Load a more powerful YOLOv8 model
 model.to(device)
 model.eval()
 
 # Load COCO class labels
-CLASS_NAMES = model.names  # YOLOv5's built-in class names
+CLASS_NAMES = model.names  # YOLO's built-in class names
 
 def preprocess_image(image):
     image = Image.fromarray(image)
@@ -21,30 +21,34 @@ def preprocess_image(image):
 
 def detect_objects(image):
     image = preprocess_image(image)
-    results = model.predict(image)  # Run YOLOv5 inference
+    results = model.predict(image)  # Run YOLO inference
 
     # Convert results to bounding box format
     image = np.array(image)
     for result in results:
-        for box, cls in zip(result.boxes.xyxy, result.boxes.cls):
+        for box, cls, conf in zip(result.boxes.xyxy, result.boxes.cls, result.boxes.conf):
             x1, y1, x2, y2 = map(int, box[:4])
             class_name = CLASS_NAMES[int(cls)]  # Get class name
+            confidence = conf.item() * 100  # Convert confidence to percentage
 
-            # Draw bounding box
-            cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
+            # Draw a bolder bounding box
+            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 4)  # Increased thickness
 
-            # Put class label
-            cv2.putText(image, class_name, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 
-                        0.5, (255, 0, 0), 2, cv2.LINE_AA)
+            # Larger text for class label
+            label = f"{class_name} ({confidence:.1f}%)"
+            cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 
+                        1, (0, 255, 0), 3, cv2.LINE_AA)  # Larger text
 
     return image
 
-# Gradio UI
+# Gradio UI with Submit button
 iface = gr.Interface(
     fn=detect_objects,
-    inputs=gr.Image(type="numpy"),
-    outputs=gr.Image(type="numpy"),
-    live=True,
+    inputs=gr.Image(type="numpy", label="Upload Image"),
+    outputs=gr.Image(type="numpy", label="Detected Objects"),
+    title="Object Detection",
+    description="Use webcam or Upload an image to detect objects.",
+    allow_flagging="never"  # Disables unwanted flags
 )
 
-iface.launch()
+iface.launch()