Spaces:
Sleeping
Sleeping
File size: 4,328 Bytes
8719e8b 2aa62cc 8719e8b 2aa62cc 8719e8b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 |
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)
|