Spaces:
Running
Running
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\nUpload 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() |