Update yolo_model.py

yolo_model.py

@@ -3,59 +3,106 @@ from ultralytics.utils.plotting import Annotator
 import numpy as np
 import cv2
 import gradio as gr
-import os
 import yolov9
 
-model = yolov9.load('Organ_detection.pt', device="cpu")
-model.conf = 0.40
-model.iou = 0.45
+# Load the first YOLOv9 model
+model1 = yolov9.load('Organ_detection.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('update_best.pt', device="cpu")
+model2.conf = 0.40
+model2.iou = 0.45
+
+def remove_lines(img):
+    # Convert the image to grayscale
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    
+    # Apply edge detection
+    edges = cv2.Canny(gray, 50, 150, apertureSize=3)
+    
+    # Detect lines using Hough Transform
+    lines = cv2.HoughLinesP(edges, 1, np.pi/180, threshold=100, minLineLength=100, maxLineGap=10)
+    
+    if lines is not None:
+        for line in lines:
+            for x1, y1, x2, y2 in line:
+                cv2.line(img, (x1, y1), (x2, y2), (255, 255, 255), 2)
+    
+    return img
 
 def Predict(img):
     objects_name = []
     cropped_images = []
-    img_name_list=[]
-    results = model(img, size=640)
-    annotator = Annotator(img, line_width=2, example=str('Organ'))
+    img_name_list = []
 
-    detections = {}
+    # Make a copy of the image for cropping
+    img_for_cropping = img.copy()
 
-    for result in results.xyxy[0]:
+    # 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 = results.names[int(class_id)]
+        label = results1.names[int(class_id)]
         confidence = float(confidence)
 
-        if label not in detections or detections[label]['confidence'] < confidence:
-            detections[label] = {
+        if label not in detections1 or detections1[label]['confidence'] < confidence:
+            detections1[label] = {
                 'box': [xmin, ymin, xmax, ymax],
                 'confidence': confidence
             }
 
-    for label, data in detections.items():
+    # Run inference using the second model
+    results2 = model2(img, size=224)
+    annotator2 = Annotator(img, line_width=2, example=str('Organ'))
+
+    detections2 = {}
+    for result in results2.xyxy[0]:
+        xmin, ymin, xmax, ymax, confidence, class_id = result
+        label = results2.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']
 
-        # Cropping the detected object
-        cropped_img = img[int(ymin):int(ymax), int(xmin):int(xmax)]
-        cropped_images.append((label, confidence, cropped_img))
+        # 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)
 
-        # Annotating the image
-        annotator.box_label([xmin, ymin, xmax, ymax], f"{label} {confidence:.2f}", color=(255, 0, 0))
+        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, cv2.COLOR_BGR2RGB)
+        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)
 
-    annotated_img = annotator.result()
-    #print(img_name_list)
-    objects_name = [(label, data['confidence']) for label, data in detections.items()]
+        # 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]
-    #print(len(labels))
-    #return annotated_img, cropped_images, objects_name
-    return img_name_list,labels
+
+    return img_name_list, labels
 
 def output_display(img):
     annotated_img, cropped_images, objects_name = Predict(img)
@@ -64,4 +111,5 @@ def output_display(img):
     crops = [crop for _, _, crop in cropped_images]
     labels = [{"label": label, "confidence": confidence} for label, confidence in objects_name]
 
-    return annotated_img, crops, labels
+    return annotated_img, crops, labels
+