Spaces:
Runtime error
Runtime error
File size: 5,820 Bytes
8f31dbb d9df1e2 8d441ae 3316ae6 d9df1e2 feb9d7d d9df1e2 8f31dbb d9df1e2 8f31dbb 830c6a4 8f31dbb eef5e41 8f31dbb eef5e41 8f31dbb eef5e41 8f31dbb eef5e41 8f31dbb 0f53572 8f31dbb f8c921c 8f31dbb 2cf60f0 8f31dbb 918bea5 7c1e71d 8f31dbb a624832 8f31dbb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 |
# Ultralytics
from ultralytics import YOLO
import torch
# Gradio
import gradio as gr
import moviepy.editor as moviepy
# System and files
import os
import glob
import uuid
# Image manipulation
import numpy as np
import cv2
print(torch.__version__)
# Use GPU if available
if torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
os.system("nvidia-smi")
print("[INFO]: Imported modules!")
track_model = YOLO('yolov8n.pt') # Load an official Detect model
print("[INFO]: Downloaded models!")
def check_extension(video):
video = os.path.join(video)
clip = moviepy.VideoFileClip(video)
split_tup = os.path.splitext(video)
print(split_tup)
# extract the file name and extension
file_name = split_tup[0]
file_extension = split_tup[1]
if file_extension != ".mp4":
print("Converting to mp4")
video = file_name+".mp4"
clip.write_videofile(video, threads = 8)
return video
def tracking(video, model, boxes=True):
print("[INFO] Is cuda available? ", torch.cuda.is_available())
print(device)
print("[INFO] Loading model...")
# Load an official or custom model
# Perform tracking with the model
print("[INFO] Starting tracking!")
# https://docs.ultralytics.com/modes/predict/
annotated_frame = model(video, boxes=boxes, device=device)
return annotated_frame
def show_tracking(video_content):
video = check_extension(video_content)
# https://docs.ultralytics.com/datasets/detect/coco/
video = cv2.VideoCapture(video_content)
# Track
video_track = tracking(video_content, track_model.track)
# Prepare to save video
#out_file = os.path.join(vis_out_dir, "track.mp4")
out_file = "track.mp4"
print("[INFO]: TRACK", out_file)
fourcc = cv2.VideoWriter_fourcc(*"mp4v") # Codec for MP4 video
fps = video.get(cv2.CAP_PROP_FPS)
height, width, _ = video_track[0][0].orig_img.shape
size = (width,height)
out_track = cv2.VideoWriter(out_file, fourcc, fps, size)
# Go through frames and write them
for frame_track in video_track:
result_track = frame_track[0].plot() # plot a BGR numpy array of predictions
out_track.write(result_track)
print("[INFO] Done with frames")
#print(type(result_pose)) numpy ndarray
out_track.release()
video.release()
cv2.destroyAllWindows() # Closing window
return out_file
def track_blocks(video_content):
files = []
for v in video_content:
files.append(show_tracking(v))
return files
block = gr.Blocks()
with block:
with gr.Column():
with gr.Tab("Record video with webcam"):
with gr.Column():
with gr.Row():
with gr.Column():
webcam_input = gr.Video(source="webcam", height=256)
with gr.Row():
submit_detect_web = gr.Button("Detect and track objects", variant="primary")
with gr.Row():
webcam_output4 = gr.Video(height=716, label = "Detection and tracking", show_label=True, format="mp4")
# with gr.Tab("Upload video"):
# with gr.Column():
# with gr.Row():
# with gr.Column():
# video_input = gr.Video(source="upload", height=256)
# with gr.Row():
# submit_detect_file = gr.Button("Detect and track objects", variant="primary")
# with gr.Row():
# video_output4 = gr.Video(height=512, label = "Detection and tracking", show_label=True, format="mp4")
with gr.Tab("General information"):
gr.Markdown("""
\n # Information about the models
\n ## Detection and tracking:
\n The tracking method in the Ultralight's YOLOv8 model is used for object tracking in videos. It takes a video file or a camera stream as input and returns the tracked objects in each frame. The method uses the COCO dataset classes for object detection and tracking.
\n The COCO dataset contains 80 classes of objects such as person, car, bicycle, etc. See https://docs.ultralytics.com/datasets/detect/coco/ for all available classes. The tracking method uses the COCO classes to detect and track the objects in the video frames. The tracked objects are represented as bounding boxes with labels indicating the class of the object.""")
# From file
#submit_detect_file.click(fn=track_blocks,
# inputs= video_input,
# outputs = video_output4,
# queue=True)
submit_detect_web.click(fn=show_tracking,
inputs= webcam_input,
outputs = webcam_output4,
queue=True)
if __name__ == "__main__":
block.queue(#concurrency_count=20, # When you increase the concurrency_count parameter in queue(), max_threads() in launch() is automatically increased as well.
max_size=30, # Maximum number of requests that the queue processes
api_open = False # When creating a Gradio demo, you may want to restrict all traffic to happen through the user interface as opposed to the programmatic API that is automatically created for your Gradio demo.
).launch(auth=("novouser", "bstad2023"))
|