Added vertical only mode

app.py

@@ -141,9 +141,10 @@ def read_motion_csv(csv_filename):
     print("[INFO] Completed reading motion CSV.")
     return motion_data
 
-def stabilize_video_using_csv(video_file, csv_file, zoom=1.0, output_file=None):
+def stabilize_video_using_csv(video_file, csv_file, zoom=1.0, vertical_only=False, output_file=None):
     """
     Stabilizes the input video using motion data from the CSV file.
+    If vertical_only is True, only vertical motion is corrected (horizontal displacement is ignored).
     """
     start_time = time.time()
     print(f"[INFO] Starting stabilization using CSV: {csv_file}")
@@ -183,6 +184,9 @@ def stabilize_video_using_csv(video_file, csv_file, zoom=1.0, output_file=None):
             frame = zoomed_frame[start_y:start_y+height, start_x:start_x+width]
         
         dx, dy = motion_data.get(frame_idx, (0, 0))
+        if vertical_only:
+            # If only vertical stabilization is desired, ignore horizontal motion.
+            dx = 0
         transform = np.array([[1, 0, dx],
                               [0, 1, dy]], dtype=np.float32)
         stabilized_frame = cv2.warpAffine(frame, transform, (width, height))
@@ -198,11 +202,12 @@ def stabilize_video_using_csv(video_file, csv_file, zoom=1.0, output_file=None):
     print(f"[INFO] Stabilized video saved to: {output_file} in {elapsed:.2f} seconds")
     return output_file
 
-def process_video_ai(video_file, zoom):
+def process_video_ai(video_file, zoom, vertical_only):
     """
     Gradio interface function:
       - Generates motion data (CSV) from the input video using an AI model (RAFT if available, else Farneback).
       - Stabilizes the video based on the generated motion data.
+      - If vertical_only is True, only vertical stabilization is applied.
     
     Returns:
         Tuple containing the original video file path, the stabilized video file path, and log output.
@@ -216,7 +221,7 @@ def process_video_ai(video_file, zoom):
     
         print("[INFO] Starting AI-powered video processing...")
         csv_file = generate_motion_csv(video_file)
-        stabilized_path = stabilize_video_using_csv(video_file, csv_file, zoom=zoom)
+        stabilized_path = stabilize_video_using_csv(video_file, csv_file, zoom=zoom, vertical_only=vertical_only)
         print("[INFO] Video processing complete.")
     logs = log_buffer.getvalue()
     return video_file, stabilized_path, logs
@@ -224,12 +229,13 @@ def process_video_ai(video_file, zoom):
 # Build the Gradio UI.
 with gr.Blocks() as demo:
     gr.Markdown("# AI-Powered Video Stabilization")
-    gr.Markdown("Upload a video and select a zoom factor. The system will generate motion data using an AI model (RAFT if available) and then stabilize the video.")
+    gr.Markdown("Upload a video, select a zoom factor, and choose whether to apply only vertical stabilization. The system will generate motion data using an AI model (RAFT if available) and then stabilize the video accordingly.")
     
     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")
+            vertical_checkbox = gr.Checkbox(label="Vertical Stabilization Only", value=False)
             process_button = gr.Button("Process Video")
         with gr.Column():
             original_video = gr.Video(label="Original Video")
@@ -238,7 +244,7 @@ with gr.Blocks() as demo:
     
     process_button.click(
         fn=process_video_ai,
-        inputs=[video_input, zoom_slider],
+        inputs=[video_input, zoom_slider, vertical_checkbox],
         outputs=[original_video, stabilized_video, logs_output]
     )