Spaces:
Sleeping
Sleeping
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='rs01313.jpg', 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() | |