app.py

# 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):
    split_tup = os.path.splitext(video)

    # extract the file name and extension
    file_name = split_tup[0]
    file_extension = split_tup[1]

    if file_extension != ".mp4":
        print("Converting to mp4")
        clip = moviepy.VideoFileClip(video)

        video = file_name+".mp4"
        clip.write_videofile(video)
    
    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):

        # 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

block = gr.Blocks()
with block:
    with gr.Column():       
        with gr.Tab("Upload video"):
            with gr.Column():
                with gr.Row():
                    with gr.Column():
                        video_input = gr.Video(source="upload", type="filepath", height=612)
                        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("Record video with webcam"):
            
            with gr.Column():
                with gr.Row():
                    with gr.Column():
                        webcam_input = gr.Video(source="webcam", height=612)
                        
                        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.""")
            gr.Markdown("You can load the keypoints in python in the following way: ")


        # From file        
        submit_detect_file.click(fn=show_tracking, 
                                inputs= video_input, 
                                outputs = video_output4,
                                queue=False)
        
    
        submit_detect_web.click(fn=show_tracking, 
                                inputs= webcam_input, 
                                outputs = webcam_output4,
                                queue=False)


if __name__ == "__main__":
    block.queue(#concurrency_count=5, # When you increase the concurrency_count parameter in queue(), max_threads() in launch() is automatically increased as well.
                #max_size=25, # 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()