Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python | |
| # coding: utf-8 | |
| # In[ ]: | |
| # import required libraries | |
| from ultralytics import YOLO | |
| import gradio as gr | |
| import cv2 | |
| import math | |
| from src.items import classNames | |
| # In[ ]: | |
| # detection function | |
| def yolo_detect(feed, vid): | |
| video = vid | |
| # Load a pretrained YOLOv8n model | |
| model = YOLO('yolov8n.pt') | |
| # Run inference on the source | |
| results = model(video, stream=True, verbose=False) | |
| frames = list() | |
| # plot annotations | |
| for frame in results: | |
| boxes = frame.boxes | |
| single = frame.orig_img | |
| for box in boxes: | |
| # bounding box | |
| x1, y1, x2, y2 = box.xyxy[0] | |
| x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) # convert to int values | |
| # put box in cam | |
| cv2.rectangle(single, (x1, y1), (x2, y2), (255, 0, 255), 3) | |
| # object details | |
| cv2.putText(single, classNames[int(box.cls[0])], (x1,y1), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1) | |
| frames.append(single) | |
| cv2.destroyAllWindows() | |
| h, w, c = frames[1].shape | |
| out_file = "output.avi" | |
| fourcc=cv2.VideoWriter_fourcc('X', 'V', 'I', 'D') | |
| writer = out = cv2.VideoWriter(out_file, fourcc, 25.0, (w, h)) | |
| for i in range(len(frames)): | |
| writer.write(frames[i]) | |
| writer.release() | |
| return out_file | |
| # In[ ]: | |
| demo = gr.Interface(fn=yolo_detect, | |
| inputs=[gr.PlayableVideo(source='webcam'), gr.Video(autoplay=True)], | |
| outputs=[gr.PlayableVideo(autoplay=True, format='avi')], | |
| cache_examples=True, allow_flagging='never') | |
| demo.queue() | |
| demo.launch(inline=False, debug=True, show_api=False, quiet=True) | |