Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| from lstr import LSTR | |
| model_path = "models/model_float32.onnx" | |
| title = "Lane Shape Prediction with Transformers (LSTR)" | |
| description = "Demo for performing lane detection using the LSTR model. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below." | |
| article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2011.04233'>End-to-end Lane Shape Prediction with Transformers</a> | <a href='https://github.com/liuruijin17/LSTR'>Original Model</a></p>" | |
| # Initialize lane detection model | |
| lane_detector = LSTR(model_path) | |
| def inference(image): | |
| image = np.array(image, dtype=np.uint8) | |
| input_img = image.copy() | |
| detected_lanes, lane_ids = lane_detector.detect_lanes(input_img) | |
| output_img = lane_detector.draw_lanes(image) | |
| # output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2RGB) | |
| im_pil = Image.fromarray(output_img) | |
| return im_pil | |
| gr.Interface( | |
| inference, | |
| [gr.inputs.Image(type="pil", label="Input")], | |
| gr.outputs.Image(type="pil", label="Output"), | |
| title=title, | |
| description=description, | |
| article=article, | |
| examples=[ | |
| ["dog_road.jpg"], | |
| ["swiss_road.jpeg"] | |
| ]).launch() | |