Update app.py

app.py

@@ -1,15 +1,21 @@
 import gradio as gr
 import cv2
 import os
-import subprocess
+import pandas as pd
+import numpy as np
+import torch
 from ultralytics import YOLO
 from ultralytics.solutions import object_counter
+import subprocess
+import spaces  # Import spaces for ZeroGPU integration
 
 # Initialize the YOLO model
 MODEL = "yolov8n.pt"
 model = YOLO(MODEL)
 model.fuse()
 
+dict_classes = model.model.names
+
 # Auxiliary functions
 def resize_frame(frame, scale_percent):
     width = int(frame.shape[1] * scale_percent / 100)
@@ -18,38 +24,129 @@ def resize_frame(frame, scale_percent):
     resized = cv2.resize(frame, dim, interpolation=cv2.INTER_AREA)
     return resized
 
-@gr.spaces.GPU
+@spaces.GPU
 def process_video(video_file, scale_percent, line_start_x, line_start_y, line_end_x, line_end_y, line_thickness, draw_tracks, view_img, view_in_counts, view_out_counts, track_thickness, region_thickness, line_dist_thresh, persist, conf, iou, classes, verbose):
-    # Video processing logic here
+    # Ensure classes is a list of integers
+    classes = [int(x) for x in classes.split(',') if x.strip().isdigit()] if classes else None
+    
+    line_points = [(line_start_x, line_start_y), (line_end_x, line_end_y)]
+    
+    cap = cv2.VideoCapture(video_file)
+    if not cap.isOpened():
+        raise ValueError("Failed to open video file")
+    
+    tmp_output_path = "processed_output_temp.mp4"
+    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) * scale_percent / 100)
+    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) * scale_percent / 100)
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    video_writer = cv2.VideoWriter(tmp_output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
+
+    counter = object_counter.ObjectCounter(
+        classes_names=model.names,
+        view_img=view_img,
+        reg_pts=line_points,
+        draw_tracks=draw_tracks,
+        line_thickness=int(line_thickness),
+        track_thickness=int(track_thickness),
+        region_thickness=int(region_thickness),
+        line_dist_thresh=line_dist_thresh,
+        view_in_counts=view_in_counts,
+        view_out_counts=view_out_counts,
+        count_reg_color=(255, 0, 255),  # Magenta
+        track_color=(0, 255, 0),  # Green
+        count_txt_color=(255, 255, 255),  # White
+        count_bg_color=(50, 50, 50)  # Dark gray
+    )
+
+    prev_frame = None
+    prev_keypoints = None
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frame = resize_frame(frame, scale_percent)
+
+        # Adjust line points based on scaling
+        scaled_line_points = [(int(x * scale_percent / 100), int(y * scale_percent / 100)) for x, y in line_points]
+        for point1, point2 in zip(scaled_line_points[:-1], scaled_line_points[1:]):
+            cv2.line(frame, tuple(map(int, point1)), tuple(map(int, point2)), (255, 255, 0), int(line_thickness))
+
+        tracks = model.track(frame, persist=persist, conf=conf, iou=iou, classes=classes, verbose=verbose)
+        
+        # Update the counter with the current frame and tracks
+        frame = counter.start_counting(frame, tracks)
+
+        # Check if the previous frame is initialized for optical flow calculation
+        if prev_frame is not None:
+            try:
+                prev_frame_resized = resize_frame(prev_frame, scale_percent)
+                matched_keypoints, status, _ = cv2.calcOpticalFlowPyrLK(prev_frame_resized, frame, prev_keypoints, None)
+                prev_keypoints = matched_keypoints
+            except cv2.error as e:
+                print(f"Error in optical flow calculation: {e}")
+
+        prev_frame = frame.copy()
+        prev_keypoints = cv2.goodFeaturesToTrack(cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY), maxCorners=100, qualityLevel=0.3, minDistance=7, blockSize=7)
+
+        video_writer.write(frame)
+
+    cap.release()
+    video_writer.release()
+
+    # Reduce the resolution of the video for download
+    output_path = "processed_output.mp4"
+    if h > 1080:
+        resolution = "1920x1080"
+    else:
+        resolution = "1280x720"
+
+    subprocess.run(
+        ["ffmpeg", "-y", "-i", tmp_output_path, "-vf", f"scale={resolution}", "-crf", "18", "-preset", "veryfast", "-hide_banner", "-loglevel", "error", output_path]
+    )
+    os.remove(tmp_output_path)
+
+    return output_path
 
 def preview_line(video_file, scale_percent, line_start_x, line_start_y, line_end_x, line_end_y, line_thickness):
-    # Generate a preview of the line on the video
-    # Code here
+    cap = cv2.VideoCapture(video_file)
+    ret, frame = cap.read()
+    if not ret:
+        raise ValueError("Failed to read video frame")
+
+    frame = resize_frame(frame, scale_percent)
+    line_points = [(line_start_x, line_start_y), (line_end_x, line_end_y)]
+    scaled_line_points = [(int(x * scale_percent / 100), int(y * scale_percent / 100)) for x, y in line_points]
+    for point1, point2 in zip(scaled_line_points[:-1], scaled_line_points[1:]):
+        cv2.line(frame, tuple(map(int, point1)), tuple(map(int, point2)), (255, 255, 0), int(line_thickness))
+    
+    preview_path = "preview_line.jpg"
+    cv2.imwrite(preview_path, frame)
+    return preview_path
 
 def gradio_app(video, scale_percent, line_start_x, line_start_y, line_end_x, line_end_y, line_thickness, draw_tracks, view_img, view_in_counts, view_out_counts, track_thickness, region_thickness, line_dist_thresh, persist, conf, iou, classes_to_track, verbose):
-    # Main Gradio app logic
-    # Code here
+    output_path = process_video(video.name, scale_percent, line_start_x, line_start_y, line_end_x, line_end_y, int(line_thickness), draw_tracks, view_img, view_in_counts, view_out_counts, int(track_thickness), int(region_thickness), line_dist_thresh, persist, conf, iou, classes_to_track, verbose)
+    return output_path, output_path
 
 def update_preview(video, scale_percent, line_start_x, line_start_y, line_end_x, line_end_y, line_thickness):
-    # Update preview logic
-    # Code here
+    return preview_line(video.name, scale_percent, line_start_x, line_start_y, line_end_x, line_end_y, int(line_thickness))
 
 def set_4k_coordinates():
     return 0, 1500, 3840, 1500
 
 def set_1080p_coordinates():
     return 0, 700, 1920, 700
-
 with gr.Blocks(css="style.css", theme="dark") as demo:
     with gr.Row():
         with gr.Column(scale=1):
             video_input = gr.File(label="Upload your video")
-            set_4k_button = gr.Button("Set 4K Line")
-            set_1080p_button = gr.Button("Set 1080p Line")
-            line_start_x = gr.Number(label="Line Start X", value=500)
-            line_start_y = gr.Number(label="Line Start Y", value=1500)
-            line_end_x = gr.Number(label="Line End X", value=3400)
-            line_end_y = gr.Number(label="Line End Y", value=1500)
+            with gr.Row():
+                set_4k_button = gr.Button("4K")
+                set_1080p_button = gr.Button("1080p")
+            line_start_x = gr.Number(label="Line Start X", value=500, precision=0)
+            line_start_y = gr.Number(label="Line Start Y", value=1500, precision=0)
+            line_end_x = gr.Number(label="Line End X", value=3400, precision=0)
+            line_end_y = gr.Number(label="Line End Y", value=1500, precision=0)
             line_thickness = gr.Slider(minimum=1, maximum=10, value=2, label="Line Thickness")
             draw_tracks = gr.Checkbox(label="Draw Tracks", value=True)
             view_img = gr.Checkbox(label="Display Image with Annotations", value=True)
@@ -61,7 +158,7 @@ with gr.Blocks(css="style.css", theme="dark") as demo:
             persist = gr.Checkbox(label="Persist Tracks", value=True)
             conf = gr.Slider(minimum=0.0, maximum=1.0, value=0.1, step=0.05, label="Confidence Threshold")
             iou = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.05, label="IOU Threshold")
-            classes_to_track = gr.Textbox(label="Classes to Track", value="2,3,5,7")
+            classes_to_track = gr.Textbox(label="Classes to Track (comma-separated ids)", value="2,3,5,7")
             verbose = gr.Checkbox(label="Verbose Tracking", value=True)
             scale_percent = gr.Slider(minimum=10, maximum=100, value=100, step=10, label="Scale Percentage")
             process_button = gr.Button("Process Video")
@@ -70,9 +167,21 @@ with gr.Blocks(css="style.css", theme="dark") as demo:
             video_output = gr.Video(label="Processed Video")
             download_button = gr.File(label="Download Processed Video")
 
-    set_4k_button.click(set_4k_coordinates, outputs=[line_start_x, line_start_y, line_end_x, line_end_y])
-    set_1080p_button.click(set_1080p_coordinates, outputs=[line_start_x, line_start_y, line_end_x, line_end_y])
+    def update_preview_and_display(video, scale_percent, line_start_x, line_start_y, line_end_x, line_end_y, line_thickness):
+        preview_path = update_preview(video, scale_percent, line_start_x, line_start_y, line_end_x, line_end_y, line_thickness)
+        return preview_path
+
+    video_input.change(update_preview_and_display, inputs=[video_input, scale_percent, line_start_x, line_start_y, line_end_x, line_end_y, line_thickness], outputs=preview_image)
+    for component in [scale_percent, line_start_x, line_start_y, line_end_x, line_end_y, line_thickness, draw_tracks, view_img, view_in_counts, view_out_counts, track_thickness, region_thickness, line_dist_thresh, persist, conf, iou, classes_to_track, verbose]:
+        component.change(update_preview_and_display, inputs=[video_input, scale_percent, line_start_x, line_start_y, line_end_x, line_end_y, line_thickness], outputs=preview_image)
+
+    set_4k_button.click(lambda: set_4k_coordinates(), outputs=[line_start_x, line_start_y, line_end_x, line_end_y])
+    set_1080p_button.click(lambda: set_1080p_coordinates(), outputs=[line_start_x, line_start_y, line_end_x, line_end_y])
+
+    def clear_previous_video():
+        return None, None
+
+    process_button.click(clear_previous_video, outputs=[video_output, download_button], queue=False)
     process_button.click(gradio_app, inputs=[video_input, scale_percent, line_start_x, line_start_y, line_end_x, line_end_y, line_thickness, draw_tracks, view_img, view_in_counts, view_out_counts, track_thickness, region_thickness, line_dist_thresh, persist, conf, iou, classes_to_track, verbose], outputs=[video_output, download_button])
-    video_input.change(update_preview, inputs=[video_input, scale_percent, line_start_x, line_start_y, line_end_x, line_end_y, line_thickness], outputs=[preview_image])
 
-demo.launch()
+demo.launch()