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)