Update app.py

app.py

@@ -6,21 +6,70 @@ import gradio as gr
 import os
 import yolov9
 
-model=yolov9.load('Organ_detection.pt',device="cpu")
+model = yolov9.load('Organ_detection.pt', device="cpu")
 model.conf = 0.40
 model.iou = 0.45
 
 def Predict(img):
-    
     objects_name = []
     cropped_images = []
-    # Object detection
+    img_name_list=[]
     results = model(img, size=640)
-    print(results)
-    output = results.render()
-    return output[0]
+    annotator = Annotator(img, line_width=2, example=str('Organ'))
 
-interface=gr.Interface(fn=Predict,inputs=["image"],
-                  outputs=["image"])
-interface.launch(debug=True)
+    detections = {}
+
+    for result in results.xyxy[0]:
+        xmin, ymin, xmax, ymax, confidence, class_id = result
+        label = results.names[int(class_id)]
+        confidence = float(confidence)
+
+        if label not in detections or detections[label]['confidence'] < confidence:
+            detections[label] = {
+                'box': [xmin, ymin, xmax, ymax],
+                'confidence': confidence
+            }
+
+    for label, data in detections.items():
+        xmin, ymin, xmax, ymax = data['box']
+        confidence = data['confidence']
+
+        # Cropping the detected object
+        cropped_img = img[int(ymin):int(ymax), int(xmin):int(xmax)]
+        cropped_images.append((label, confidence, cropped_img))
+
+        # Annotating the image
+        annotator.box_label([xmin, ymin, xmax, ymax], f"{label} {confidence:.2f}", color=(255, 0, 0))
+
+        # Convert the cropped image from BGR to RGB before saving
+        cropped_img_rgb = cv2.cvtColor(cropped_img, cv2.COLOR_BGR2RGB)
 
+        # Save the cropped image
+        crop_filename = f"{label}.jpg"
+        img_name_list.append(crop_filename)
+        cv2.imwrite(crop_filename, cropped_img_rgb)
+
+    annotated_img = annotator.result()
+    #print(img_name_list)
+    objects_name = [(label, data['confidence']) for label, data in detections.items()]
+    labels = [{"label": label, "confidence": confidence} for label, confidence in objects_name]
+    #print(len(labels))
+    return annotated_img, cropped_images, objects_name
+    #return img_name_list,labels
+
+def output_display(img):
+    annotated_img, cropped_images, objects_name = Predict(img)
+
+    # Extract cropped images and labels separately
+    crops = [crop for _, _, crop in cropped_images]
+    labels = [{"label": label, "confidence": confidence} for label, confidence in objects_name]
+
+    return annotated_img, crops, labels
+
+interface = gr.Interface(fn=output_display,
+                         inputs=["image"],
+                         outputs=[gr.Image(label="Annotated Image"),
+                                  gr.Gallery(label="Cropped Images"),
+                                  gr.JSON(label="Labels and Confidence")])
+
+interface.launch(debug=True)