Rehman1603's picture
Update app.py
52d2429 verified
raw
history blame
1.58 kB
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)