Spaces:
Sleeping
Sleeping
| from ultralytics import YOLO | |
| 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 | |
| def Predict(img): | |
| objects_name = [] | |
| cropped_images = [] | |
| img_name_list=[] | |
| results = model(img, size=640) | |
| annotator = Annotator(img, line_width=2, example=str('Organ')) | |
| 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) | |