Update app.py

app.py

@@ -20,30 +20,22 @@ class_names = {0: "Dumbbell", 1: "Kettlebell"}  # Adjust based on your dataset
 
 # Define prediction function
 def predict(image):
-    # Perform inference with a higher confidence threshold
-    results = model.predict(image, conf=0.5)  # Set confidence threshold
-
-    # Annotate the output image
-    output_image = results[0].plot()
-
-    # Filter and prepare detection details
-    detections = []
-    for result in results[0].boxes:
-        class_idx = int(result.cls)
-        confidence = float(result.conf)
-        bbox = result.xyxy.tolist()
-
-        # Restrict to trained classes
-        if class_idx < len(class_names):  # Ensure index is valid
-            detections.append({
-                "class": class_names[class_idx],  # Use the mapped class name
-                "confidence": confidence,
-                "bbox": bbox
-            })
+    # Perform inference
+    results = model.predict(image)
+    output_image = results[0].plot()  # Annotated image with predictions
+
+    # Prepare detection details
+    detections = [
+        {
+            "class": class_names.get(int(result.cls), "Unknown"),  # Map class index to name
+            "confidence": float(result.conf),
+            "bbox": result.xyxy.tolist(),
+        }
+        for result in results[0].boxes
+    ]
 
     return output_image, detections
 
-
 # Create Gradio interface
 interface = gr.Interface(
     fn=predict,