import gradio as gr
from TheDistanceAssessor import run, load_segformer, load_yolo

def process_image(input_image):
    image_size = [1024,1024]
    target_distances = [650,1000,2000]
    num_ys = 10

    PATH_model_seg = 'SegFormer_B3_1024_finetuned.pth'
    PATH_model_det = 'yolov8s.pt'
    model_seg = load_segformer(PATH_model_seg)
    model_det = load_yolo(PATH_model_det)
    output_image = run(input_image, model_seg, model_det, image_size, target_distances, num_ys = num_ys)
    return output_image

# Create the Gradio interface
iface = gr.Interface(
    fn=process_image,  # The function to be called
    inputs=gr.Image(type="pil"),  # Input type
    outputs=gr.Image(type="numpy"),  # Output type
    title="Image Processor",  # Title of the interface
    description="Upload an image and get a processed image as output."  # Description of the interface
)

# Launch the interface
if __name__ == "__main__":
    iface.launch()