Update app.py

app.py

@@ -57,7 +57,6 @@ def compress_video(video_file, target_width, target_height, progress=gr.Progress
         ret, frame = cap.read()
         if not ret:
             break
-        # Resize frame to target resolution
         compressed_frame = cv2.resize(frame, (new_width, new_height))
         out.write(compressed_frame)
         if frame_idx % 10 == 0 or frame_idx == total_frames:
@@ -129,12 +128,10 @@ def generate_motion_csv(video_file, output_csv=None, progress=gr.Progress(), pro
                                                     iterations=3, poly_n=5, poly_sigma=1.2, flags=0)
                 prev_gray = curr_gray
 
-            # Compute median magnitude and angle.
             mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1], angleInDegrees=True)
             median_mag = np.median(mag)
             median_ang = np.median(ang)
             
-            # Compute "zoom factor": fraction of pixels moving away from center.
             h, w = flow.shape[:2]
             center_x, center_y = w / 2, h / 2
             x_coords, y_coords = np.meshgrid(np.arange(w), np.arange(h))
@@ -223,7 +220,6 @@ def stabilize_video_using_csv(video_file, csv_file, zoom=1.0, vertical_only=Fals
         if not ret:
             break
         
-        # Optionally apply zoom (resize and center-crop)
         if zoom != 1.0:
             zoomed_frame = cv2.resize(frame, None, fx=zoom, fy=zoom, interpolation=cv2.INTER_LINEAR)
             zoomed_h, zoomed_w = zoomed_frame.shape[:2]
@@ -236,7 +232,6 @@ def stabilize_video_using_csv(video_file, csv_file, zoom=1.0, vertical_only=Fals
             dx = 0  # Only vertical stabilization.
         transform = np.array([[1, 0, dx],
                               [0, 1, dy]], dtype=np.float32)
-        # Use BORDER_REPLICATE to avoid black borders
         stabilized_frame = cv2.warpAffine(frame, transform, (width, height), borderMode=cv2.BORDER_REPLICATE)
         
         out.write(stabilized_frame)
@@ -252,11 +247,12 @@ def stabilize_video_using_csv(video_file, csv_file, zoom=1.0, vertical_only=Fals
     print(f"[INFO] Stabilized video saved to: {output_file} in {elapsed:.2f} seconds")
     return output_file
 
-def process_video_ai(video_file, zoom, vertical_only, compress_mode, target_width, target_height, progress=gr.Progress(track_tqdm=True)):
+def process_video_ai(video_file, zoom, vertical_only, compress_mode, target_width, target_height, auto_zoom, progress=gr.Progress(track_tqdm=True)):
     """
     Gradio interface function:
       - Optionally compresses the video if compress_mode is True, resizing it to the chosen resolution.
       - Generates motion data from the (possibly compressed) video.
+      - If auto_zoom is enabled, computes the optimal zoom level based on the maximum cumulative translation so that no black borders appear.
       - Stabilizes the video based on the generated motion data.
       - If vertical_only is True, only vertical stabilization is applied.
     
@@ -274,10 +270,8 @@ def process_video_ai(video_file, zoom, vertical_only, compress_mode, target_widt
         # If compression is enabled, compress the video first.
         if compress_mode:
             gr.Info("Compressing video before processing...")
-            # Compression phase uses progress 0 to 0.2
             video_file = compress_video(video_file, target_width, target_height, progress=progress, progress_offset=0.0, progress_scale=0.2)
             gr.Info("Video compression complete.")
-            # Set new progress offsets for subsequent phases.
             motion_offset = 0.2
             motion_scale = 0.4
             stabilization_offset = 0.6
@@ -290,6 +284,31 @@ def process_video_ai(video_file, zoom, vertical_only, compress_mode, target_widt
 
         csv_file = generate_motion_csv(video_file, progress=progress, progress_offset=motion_offset, progress_scale=motion_scale)
         gr.Info("Motion CSV generated successfully.")
+        
+        # If auto zoom is enabled, compute the optimal zoom factor.
+        if auto_zoom:
+            gr.Info("Auto Zoom Mode enabled. Computing optimal zoom factor...")
+            motion_data = read_motion_csv(csv_file)
+            dx_list = [abs(v[0]) for v in motion_data.values()]
+            dy_list = [abs(v[1]) for v in motion_data.values()]
+            max_dx = max(dx_list) if dx_list else 0
+            max_dy = max(dy_list) if dy_list else 0
+            cap = cv2.VideoCapture(video_file)
+            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+            cap.release()
+            if width - 2*max_dx > 0:
+                zoom_x = width / (width - 2*max_dx)
+            else:
+                zoom_x = 1.0
+            if height - 2*max_dy > 0:
+                zoom_y = height / (height - 2*max_dy)
+            else:
+                zoom_y = 1.0
+            auto_zoom_factor = max(1.0, zoom_x, zoom_y)
+            gr.Info(f"Auto zoom factor computed: {auto_zoom_factor:.2f}")
+            zoom = auto_zoom_factor
+
         stabilized_path = stabilize_video_using_csv(video_file, csv_file, zoom=zoom, vertical_only=vertical_only,
                                                       progress=progress, progress_offset=stabilization_offset, progress_scale=stabilization_scale)
         gr.Info("Video stabilization complete.")
@@ -301,7 +320,7 @@ def process_video_ai(video_file, zoom, vertical_only, compress_mode, target_widt
 with gr.Blocks() as demo:
     gr.Markdown("# AI-Powered Video Stabilization")
     gr.Markdown(
-        "Upload a video, select a zoom factor, choose whether to apply only vertical stabilization, and optionally compress the video before processing. "
+        "Upload a video, select a zoom factor (or use Auto Zoom Mode), choose whether to apply only vertical stabilization, and optionally compress the video before processing. "
         "If compressing, specify the target resolution (width and height) for the compressed video. "
         "The system will generate motion data using an AI model (RAFT if available) and then stabilize the video with live progress updates and alerts."
     )
@@ -309,7 +328,8 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             video_input = gr.Video(label="Input Video")
-            zoom_slider = gr.Slider(minimum=1.0, maximum=2.0, step=0.1, value=1.0, label="Zoom Factor")
+            zoom_slider = gr.Slider(minimum=1.0, maximum=3.0, step=0.1, value=1.0, label="Zoom Factor (ignored if Auto Zoom enabled)")
+            auto_zoom_checkbox = gr.Checkbox(label="Auto Zoom Mode", value=False)
             vertical_checkbox = gr.Checkbox(label="Vertical Stabilization Only", value=False)
             compress_checkbox = gr.Checkbox(label="Compress Video Before Processing", value=False)
             target_width = gr.Number(label="Target Width (px)", value=640)
@@ -322,7 +342,7 @@ with gr.Blocks() as demo:
     
     process_button.click(
         fn=process_video_ai,
-        inputs=[video_input, zoom_slider, vertical_checkbox, compress_checkbox, target_width, target_height],
+        inputs=[video_input, zoom_slider, vertical_checkbox, compress_checkbox, target_width, target_height, auto_zoom_checkbox],
         outputs=[original_video, stabilized_video, logs_output]
     )