import gradio as gr
import yolov7
import tempfile
import cv2
from pathlib import Path

def image_fn(image, model_path, image_size, conf_threshold, iou_threshold):
    """
    YOLOv7 inference function
    Args:
        image: Input image
        model_path: Path to the model
        image_size: Image size
        conf_threshold: Confidence threshold
        iou_threshold: IOU threshold
    Returns:
        Rendered image
    """
    model = yolov7.load(model_path, device="cpu", hf_model=True, trace=False)
    model.conf = conf_threshold
    model.iou = iou_threshold
    results = model([image], size=image_size)
    img = results.render()[0]
    img.flags.writeable = True  # Make the image writable
    return img

# Use gr.components instead of gr.inputs and gr.outputs
image_interface = gr.Interface(
    fn=image_fn,
    inputs=[
        gr.components.Image(type="pil", label="Input Image"),
        gr.components.Dropdown(
            choices=["Aalaa/Yolov7_Visual_Pollution_Detection"],
            value="Aalaa/Yolov7_Visual_Pollution_Detection",
            label="Model"
        ),
        gr.components.Slider(minimum=320, maximum=1280, value=640, step=32, label="Image Size"),
        gr.components.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.05, label="Confidence Threshold"),
        gr.components.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU Threshold")
    ],
    outputs=gr.components.Image(type="numpy", label="Output Image"),
    examples=[['image1.jpg', 'Aalaa/Yolov7_Visual_Pollution_Detection', 640, 0.25, 0.45]],
    cache_examples=True,
    theme='default'
)

if __name__ == "__main__":
    gr.TabbedInterface(
        [image_interface],
        ["Run on Images"],
    ).launch()