Update app.py

app.py

@@ -29,8 +29,15 @@ def count_frames(video_file):
     cap.release()
     return total_frames
 
+def file_size_check(file_path, max_size_mb=500):
+    if os.path.getsize(file_path) > max_size_mb * 1024 * 1024:
+        raise ValueError("File size exceeds the maximum limit of {} MB".format(max_size_mb))
+
 @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, progress):
+    # Check file size
+    file_size_check(video_file)
+    
     # Ensure classes is a list of integers
     classes = [int(x) for x in classes.split(',') if x.strip().isdigit()] if classes else None
     
@@ -68,42 +75,43 @@ def process_video(video_file, scale_percent, line_start_x, line_start_y, line_en
     prev_keypoints = None
     processed_frames = 0
 
-    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)
-        
-        # Update progress
-        processed_frames += 1
-        progress(processed_frames / total_frames)
-
-    cap.release()
-    video_writer.release()
+    try:
+        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)
+            
+            # Update progress
+            processed_frames += 1
+            progress(processed_frames / total_frames)
+    finally:
+        cap.release()
+        video_writer.release()
 
     # Reduce the resolution of the video for download
     output_path = "processed_output.mp4"
@@ -112,15 +120,24 @@ def process_video(video_file, scale_percent, line_start_x, line_start_y, line_en
     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)
+    try:
+        subprocess.run(
+            ["ffmpeg", "-y", "-i", tmp_output_path, "-vf", f"scale={resolution}", "-crf", "18", "-preset", "veryfast", "-hide_banner", "-loglevel", "error", output_path],
+            check=True
+        )
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError(f"Error during video processing: {e}")
+
+    if os.path.exists(tmp_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):
     cap = cv2.VideoCapture(video_file)
+    if not cap.isOpened():
+        raise ValueError("Failed to read video frame")
+
     ret, frame = cap.read()
     if not ret:
         raise ValueError("Failed to read video frame")
@@ -150,7 +167,7 @@ def set_4k_coordinates():
 def set_1080p_coordinates():
     return 0, 700, 1920, 700
 
-with gr.Blocks(theme="dark") as demo:
+with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column(scale=1):
             video_input = gr.File(label="Upload your video")  # Removed max_size parameter