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) | |