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()