Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import cv2 | |
| import codecs | |
| import os | |
| from PIL import Image | |
| from ultralytics import YOLO | |
| ################## MODEL ################## | |
| model = YOLO('best.pt') | |
| title = "RSUD20K: A Dataset for Road Scene Understanding In Autonomous Driving" | |
| description = codecs.open("description.html", "r", "utf-8").read() | |
| ################## IMAGE ################## | |
| Image_directory = "examples/images" | |
| inputs_image = [ | |
| gr.components.Image(type="filepath", label="Input Image"), | |
| ] | |
| outputs_image = [ | |
| gr.components.Image(type="numpy", label="Output Image"), | |
| ] | |
| def show_preds_image(image_path): | |
| image = cv2.imread(image_path) | |
| outputs = model.predict(source=image_path) | |
| results = Image.fromarray(outputs[0].plot()[:, :, ::-1]) | |
| return results | |
| demo_image = gr.Interface( | |
| fn=show_preds_image, | |
| title=title, | |
| description=description, | |
| inputs= inputs_image, | |
| outputs= outputs_image, | |
| examples= [os.path.join(Image_directory, fname) for fname in os.listdir(Image_directory) if fname.endswith(".jpg")], | |
| allow_flagging="never", | |
| analytics_enabled=False, | |
| ) | |
| ################## VIDEO ################## | |
| Video_directory = "examples/videos" | |
| inputs_video = [ | |
| gr.components.Video(label="Input Video"), | |
| ] | |
| outputs_video = [ | |
| gr.components.Image(type = "numpy", label="Output Video"), | |
| ] | |
| def show_preds_video(video_path): | |
| cap = cv2.VideoCapture(video_path) | |
| predicted_frames = [] | |
| while(cap.isOpened()): | |
| ret, frame = cap.read() | |
| if ret: | |
| frame_copy = frame.copy() | |
| outputs = model.predict(source=frame) | |
| results = Image.fromarray(outputs[0].plot()[:, :, ::-1]) | |
| yield results | |
| else: | |
| break | |
| cap.release() | |
| cv2.destroyAllWindows() | |
| demo_video = gr.Interface( | |
| fn=show_preds_video, | |
| title=title, | |
| description=description, | |
| inputs= inputs_video, | |
| outputs= outputs_video, | |
| examples= [os.path.join(Video_directory, fname) for fname in os.listdir(Video_directory) if fname.endswith(".mp4")], | |
| allow_flagging="never", | |
| analytics_enabled=False, | |
| ) | |
| ################## LAUNCH ################## | |
| gr.TabbedInterface( | |
| [demo_image, demo_video], | |
| tab_names=['Image inference', 'Video inference'] | |
| ).queue().launch() |