Update app.py

app.py

@@ -56,13 +56,14 @@ def detect_objects(frame):
                     
         return detections
 
-def display_detections(image, detections):
-    for x1, y1, x2, y2, conf, cls in detections:
-        label = f'{cls} ({conf:.2f})'
-        cv2.rectangle(image, (x1, y1), (x2, y2), colors.get(cls, (255, 255, 255)), 2)
-        cv2.putText(image, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, colors.get(cls, (255, 255, 255)), 2)
-    return image
-
+def display_detections(input_image, output_image, detections):
+    for det in detections:
+        x1, y1, x2, y2, conf, detected_name = det
+        color = colors.get(detected_name, (255, 255, 255))
+        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 2)
+        cv2.putText(output_image, f"{detected_name} ({conf:.2f})", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+    return output_image
+    
 def app():
     st.title("Helmet Detection App")
     st.write("This app uses YOLOv5 to detect helmets and riders in images and videos.")