Update app.py

app.py

@@ -5,15 +5,10 @@ import cv2
 import gradio as gr
 import yolov9
 
-# Load the first YOLOv9 model
-model1 = yolov9.load('best (1).pt', device="cpu")
-model1.conf = 0.40
-model1.iou = 0.45
-
-# Load the second YOLO model (assuming you have a second YOLOv9 model or another YOLO model)
-model2 = yolov9.load('best (1).pt', device="cpu")
-model2.conf = 0.40
-model2.iou = 0.45
+# Load the YOLOv9 model
+model = yolov9.load('best (1).pt', device="cpu")
+model.conf = 0.40
+model.iou = 0.45
 
 def remove_lines(img):
     # Convert the image to grayscale
@@ -33,91 +28,27 @@ def remove_lines(img):
     return img
 
 def Predict(img):
-    objects_name = []
-    cropped_images = []
-    img_name_list = []
-
-    # Make a copy of the image for cropping
-    img_for_cropping = img.copy()
-
-    # Run inference using the first model
-    results1 = model1(img, size=224)
-    annotator1 = Annotator(img, line_width=2, example=str('Organ'))
-
-    detections1 = {}
-    for result in results1.xyxy[0]:
-        xmin, ymin, xmax, ymax, confidence, class_id = result
-        label = results1.names[int(class_id)]
-        confidence = float(confidence)
-
-        if label not in detections1 or detections1[label]['confidence'] < confidence:
-            detections1[label] = {
-                'box': [xmin, ymin, xmax, ymax],
-                'confidence': confidence
-            }
-
-    # Run inference using the second model
-    results2 = model2(img, size=224)
-    annotator2 = Annotator(img, line_width=2, example=str('Organ'))
+    # Run inference using the model
+    results = model(img, size=224)
+    annotator = Annotator(img, line_width=2, example=str('Organ'))
 
-    detections2 = {}
-    for result in results2.xyxy[0]:
+    for result in results.xyxy[0]:
         xmin, ymin, xmax, ymax, confidence, class_id = result
-        label = results2.names[int(class_id)]
+        label = results.names[int(class_id)]
         confidence = float(confidence)
 
-        if label not in detections2 or detections2[label]['confidence'] < confidence:
-            detections2[label] = {
-                'box': [xmin, ymin, xmax, ymax],
-                'confidence': confidence
-            }
-
-    # Combine detections from both models
-    combined_detections = {**detections1, **detections2}
-
-    for label, data in combined_detections.items():
-        xmin, ymin, xmax, ymax = data['box']
-        confidence = data['confidence']
+        # Annotate the image
+        annotator.box_label([xmin, ymin, xmax, ymax], f"{label} {confidence:.2f}", color=(255, 0, 0))
 
-        # Cropping the detected object from the original image
-        cropped_img = img_for_cropping[int(ymin):int(ymax), int(xmin):int(xmax)]
-        
-        # Remove lines from the cropped image
-        cropped_img_cleaned = remove_lines(cropped_img)
-
-        cropped_images.append((label, confidence, cropped_img_cleaned))
-
-        # Convert the cropped image from BGR to RGB before saving
-        cropped_img_rgb = cv2.cvtColor(cropped_img_cleaned, 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)
-
-        # Annotating the image (after cropping to ensure the line is not in the cropped images)
-        annotator1.box_label([xmin, ymin, xmax, ymax], f"{label} {confidence:.2f}", color=(255, 0, 0))
-
-    annotated_img = annotator1.result()
-    objects_name = [(label, data['confidence']) for label, data in combined_detections.items()]
-    labels = [{"label": label, "confidence": confidence} for label, confidence in objects_name]
-
-    return annotated_img, cropped_images, objects_name
+    annotated_img = annotator.result()
+    return annotated_img
 
 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
+    annotated_img = Predict(img)
+    return annotated_img
 
 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")])
+                         outputs=gr.Image(label="Annotated Image"))
 
 interface.launch(debug=True)
-