import PIL.Image as Image
import PIL.ExifTags
import gradio as gr
from ultralytics import YOLO

# Update the path to the YOLOv8 weights file
weights_path = "best (2).pt"
model = YOLO(weights_path)

def correct_image_orientation(img):
    """Correct image orientation based on EXIF metadata."""
    try:
        exif = img._getexif()
        if exif is not None:
            orientation_tag = [key for key, value in PIL.ExifTags.TAGS.items() if value == 'Orientation']
            if orientation_tag:
                orientation = exif.get(orientation_tag[0])
                if orientation == 3:
                    img = img.rotate(180, expand=True)
                elif orientation == 6:
                    img = img.rotate(270, expand=True)
                elif orientation == 8:
                    img = img.rotate(90, expand=True)
    except (AttributeError, KeyError, IndexError):
        # If there is no EXIF data or any error occurs, just return the image as is
        pass
    return img

def predict_image(img, conf_threshold, iou_threshold):
    """Predicts and plots labeled objects in an image using YOLOv8 model with adjustable confidence and IOU thresholds. """
    # Correct the image orientation first
    img = correct_image_orientation(img)

    # Perform prediction using YOLOv8 model
    results = model.predict(
        source=img,
        conf=conf_threshold,
        iou=iou_threshold,
        show_labels=True,
        show_conf=True,
        imgsz=640,
    )

    # Convert the results back to a PIL image
    if results:
        im_array = results[0].plot()
        im = Image.fromarray(im_array)
        
        # Ensure the output image is horizontally oriented
        im = im.rotate(0, expand=True)

        return im

    return img

iface = gr.Interface(
    fn=predict_image,
    inputs=[
        gr.inputs.Image(type="pil", label="Upload Image"),
        gr.inputs.Slider(minimum=0, maximum=1, default=0.25, label="Confidence threshold"),
        gr.inputs.Slider(minimum=0, maximum=1, default=0.45, label="IoU threshold")
    ],
    outputs=gr.outputs.Image(type="pil", label="Result"),
    title="U-SUB DETECTRON",
     description="""Upload images for inference. The Ultralytics YOLOv8n model is used by default. 
    DO NOT TAKE PICTURE VERTICALLY, current bug prevents model detection due to random rotations. TAKE PICTURE HORIZONTALLY! (input integral image should have bottom of  integral facing downwards or leftwards)""",
    examples=[
        ["p1.png", 0.25, 0.45],
        ["p2.png", 0.25, 0.45],
        ["test1 (1).png", 0.25, 0.45],
        ["p3.jpeg", 0.25, 0.45],
        ["p4.jpg", 0.25, 0.45],
    ]
)

if __name__ == '__main__':
    iface.launch()