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 accesible from: https://github.com/oValach/RailSafeNet_DT \n\n
    Upload an image with a scene including rail track and get a processed image with marked rail critical areas and detected and classified objects."
)

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