Spaces:
Sleeping
Sleeping
| from ultralytics import YOLO | |
| from ultralytics.utils.plotting import Annotator | |
| import numpy as np | |
| 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 | |
| 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 = [] | |
| # 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')) | |
| 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 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 | |
| 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) | |