Update app.py

app.py

@@ -16,19 +16,32 @@ model = YOLO('best.pt')
 path = [['Cyst.jpg'], ['Stone.jpg'],['Normal.jpg']]
 
 
-def show_preds_image(image_path):
+def detect_objects_on_image(image_path):
     image = cv2.imread(image_path)
-    outputs = model.predict(source=image_path)
-    results = outputs[0].cpu().numpy()
-    for i, det in enumerate(results.boxes.xyxy):
+    model = YOLO("best.pt")
+    results = model.predict(image_path)
+    result = results[0]
+    output = []
+    for box in result.boxes:
+        x1, y1, x2, y2 = [
+            round(x) for x in box.xyxy[0].tolist()
+            
+        ]
+        class_id = box.cls[0].item()
+        prob = round(box.conf[0].item(), 2)
+        output.append([
+            x1, y1, x2, y2, result.names[class_id], prob
+        ])
         cv2.rectangle(
             image,
-            (int(det[0]), int(det[1])),
-            (int(det[2]), int(det[3])),
+            (x1, y1),
+            (x2, y2),
             color=(0, 0, 255),
             thickness=2,
             lineType=cv2.LINE_AA
         )
+        
+        cv2.putText(image,result.names[class_id],  (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 2)
     return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
 inputs_image = [
@@ -38,7 +51,7 @@ outputs_image = [
     gr.components.Image(type="numpy", label="Output Image"),
 ]
 interface_image = gr.Interface(
-    fn=show_preds_image,
+    fn=detect_objects_on_image,
     inputs=inputs_image,
     outputs=outputs_image,
     title="Kidney Stone and Cyst detection",