|
|
|
from ultralytics import YOLO |
|
import torch |
|
|
|
|
|
import gradio as gr |
|
import moviepy.editor as moviepy |
|
|
|
|
|
|
|
import os |
|
import glob |
|
import uuid |
|
|
|
|
|
import numpy as np |
|
import cv2 |
|
|
|
print(torch.__version__) |
|
|
|
|
|
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') |
|
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) |
|
|
|
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...") |
|
|
|
|
|
|
|
print("[INFO] Starting tracking!") |
|
|
|
annotated_frame = model(video, boxes=boxes, device=device) |
|
|
|
return annotated_frame |
|
|
|
def show_tracking(video_content): |
|
|
|
video = check_extension(video_content) |
|
|
|
|
|
video = cv2.VideoCapture(video_content) |
|
|
|
|
|
video_track = tracking(video_content, track_model.track) |
|
|
|
|
|
|
|
out_file = "track.mp4" |
|
print("[INFO]: TRACK", out_file) |
|
|
|
fourcc = cv2.VideoWriter_fourcc(*"mp4v") |
|
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) |
|
|
|
|
|
for frame_track in video_track: |
|
result_track = frame_track[0].plot() |
|
out_track.write(result_track) |
|
|
|
print("[INFO] Done with frames") |
|
|
|
|
|
out_track.release() |
|
|
|
video.release() |
|
cv2.destroyAllWindows() |
|
|
|
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("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.""") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
submit_detect_web.click(fn=show_tracking, |
|
inputs= webcam_input, |
|
outputs = webcam_output4, |
|
queue=True) |
|
|
|
|
|
if __name__ == "__main__": |
|
block.queue( |
|
max_size=30, |
|
api_open = False |
|
).launch(auth=("novouser", "bstad2023")) |
|
|
|
|
|
|