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)