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 YOLOv9 model | |
model = yolov9.load('best (1).pt', device="cpu") | |
model.conf = 0.40 | |
model.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): | |
# Run inference using the model | |
results = model(img, size=224) | |
annotator = Annotator(img, line_width=2, example=str('Organ')) | |
for result in results.xyxy[0]: | |
xmin, ymin, xmax, ymax, confidence, class_id = result | |
label = results.names[int(class_id)] | |
confidence = float(confidence) | |
# Annotate the image | |
annotator.box_label([xmin, ymin, xmax, ymax], f"{label} {confidence:.2f}", color=(255, 0, 0)) | |
annotated_img = annotator.result() | |
return annotated_img | |
def output_display(img): | |
annotated_img = Predict(img) | |
return annotated_img | |
interface = gr.Interface(fn=output_display, | |
inputs=["image"], | |
outputs=gr.Image(label="Annotated Image")) | |
interface.launch(debug=True) | |