import cv2
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)
    
    input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)
    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="numpy"),  # Input type
    outputs=gr.Image(type="numpy"),  # Output type
    title="RailSafeNet - Automatic Detection of Objects in the Track",  # Title of the interface
    description="This is a demo of the master's thesis focused on the Automatic Detection of Objects in the Track.\n The repository with the code is accessible from: https://github.com/oValach/RailSafeNet_DT \n\nUpload an image with a scene including rail track and get a processed image with marked rail critical areas and detected and classified objects."
)

example_images = gr.Markdown(
    """
    ## Example input
    Here are three example input images, you can drag and drop them into the image input window:
    """
)

example_image1 = gr.Image(value='test_file.png', type='numpy', label="Example Image 1", width=420, height=280)
example_image2 = gr.Image(value='rs00042.jpg', type='numpy', label="Example Image 2", width=420, height=280)
example_image3 = gr.Image(value='frame_0474.jpg', type='numpy', label="Example Image 3", width=420, height=280)

# Combine the interface and example images
app = gr.Blocks()

with app:
    iface.render()
    example_images.render()
    with gr.Row():
        example_image1.render()
        example_image2.render()
        example_image3.render()

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