Added compression mode

app.py

@@ -27,11 +27,56 @@ except Exception as e:
     raft_model = None
     gr.Warning("Falling back to OpenCV Farneback optical flow.")
 
-def generate_motion_csv(video_file, output_csv=None, progress=gr.Progress(), progress_offset=0.0, progress_scale=0.5):
+def compress_video(video_file, compression_factor, progress=gr.Progress(), progress_offset=0.0, progress_scale=0.2, output_file=None):
+    """
+    Compresses the video by resizing each frame to a lower resolution.
+    The new resolution is (original_width * compression_factor, original_height * compression_factor).
+    Updates progress from progress_offset to progress_offset+progress_scale.
+    """
+    start_time = time.time()
+    cap = cv2.VideoCapture(video_file)
+    if not cap.isOpened():
+        raise gr.Error("Could not open video file for compression.")
+    
+    original_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    original_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    new_width = max(1, int(original_width * compression_factor))
+    new_height = max(1, int(original_height * compression_factor))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    
+    if output_file is None:
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
+        output_file = temp_file.name
+        temp_file.close()
+    
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_file, fourcc, fps, (new_width, new_height))
+    
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    frame_idx = 1
+    print(f"[INFO] Starting video compression: {total_frames} frames, target resolution: {new_width}x{new_height}")
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        # Resize frame to new resolution
+        compressed_frame = cv2.resize(frame, (new_width, new_height))
+        out.write(compressed_frame)
+        if frame_idx % 10 == 0 or frame_idx == total_frames:
+            print(f"[INFO] Compressed frame {frame_idx}/{total_frames}")
+        progress(progress_offset + (frame_idx/total_frames)*progress_scale, desc="Compressing Video")
+        frame_idx += 1
+
+    cap.release()
+    out.release()
+    elapsed = time.time() - start_time
+    print(f"[INFO] Compressed video saved to: {output_file} in {elapsed:.2f} seconds")
+    return output_file
+
+def generate_motion_csv(video_file, output_csv=None, progress=gr.Progress(), progress_offset=0.0, progress_scale=0.4):
     """
     Generates a CSV file with motion data (columns: frame, mag, ang, zoom) from an input video.
     Uses RAFT if available, otherwise falls back to OpenCV's Farneback optical flow.
-    
     Updates progress from progress_offset to progress_offset+progress_scale.
     """
     start_time = time.time()
@@ -121,7 +166,6 @@ def generate_motion_csv(video_file, output_csv=None, progress=gr.Progress(), pro
 def read_motion_csv(csv_filename):
     """
     Reads a motion CSV file and computes cumulative offset per frame.
-    
     Returns a dictionary mapping frame numbers to (dx, dy) offsets.
     """
     print(f"[INFO] Reading motion CSV: {csv_filename}")
@@ -143,11 +187,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, vertical_only=False, progress=gr.Progress(), progress_offset=0.5, progress_scale=0.5, output_file=None):
+def stabilize_video_using_csv(video_file, csv_file, zoom=1.0, vertical_only=False, progress=gr.Progress(), progress_offset=0.6, progress_scale=0.4, output_file=None):
     """
-    Stabilizes the input video using motion data from the CSV.
+    Stabilizes the video using motion data from the CSV.
     If vertical_only is True, only vertical motion is corrected.
-    
     Updates progress from progress_offset to progress_offset+progress_scale.
     """
     start_time = time.time()
@@ -207,30 +250,46 @@ 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, progress=gr.Progress(track_tqdm=True)):
+def process_video_ai(video_file, zoom, vertical_only, compress_mode, compression_factor, progress=gr.Progress(track_tqdm=True)):
     """
     Gradio interface function:
-      - Generates motion data from the input video using an AI model (RAFT if available, else Farneback).
+      - Optionally compresses the video if compress_mode is True.
+      - Generates motion data from the (possibly compressed) video.
       - Stabilizes the video based on the generated motion data.
       - If vertical_only is True, only vertical stabilization is applied.
     
     Returns:
       Tuple: (original video file path, stabilized video file path, log output)
     """
-    # Display an info alert.
     gr.Info("Starting AI-powered video processing...")
-    
     log_buffer = io.StringIO()
     with redirect_stdout(log_buffer):
         if isinstance(video_file, dict):
             video_file = video_file.get("name", None)
         if video_file is None:
             raise gr.Error("Please upload a video file.")
-    
-        csv_file = generate_motion_csv(video_file, progress=progress, progress_offset=0.0, progress_scale=0.5)
+        
+        # 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, compression_factor, 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
+            stabilization_scale = 0.4
+        else:
+            motion_offset = 0.0
+            motion_scale = 0.5
+            stabilization_offset = 0.5
+            stabilization_scale = 0.5
+
+        csv_file = generate_motion_csv(video_file, progress=progress, progress_offset=motion_offset, progress_scale=motion_scale)
         gr.Info("Motion CSV generated successfully.")
         stabilized_path = stabilize_video_using_csv(video_file, csv_file, zoom=zoom, vertical_only=vertical_only,
-                                                      progress=progress, progress_offset=0.5, progress_scale=0.5)
+                                                      progress=progress, progress_offset=stabilization_offset, progress_scale=stabilization_scale)
         gr.Info("Video stabilization complete.")
         print("[INFO] Video processing complete.")
     logs = log_buffer.getvalue()
@@ -239,13 +298,18 @@ def process_video_ai(video_file, zoom, vertical_only, progress=gr.Progress(track
 # Build the Gradio UI.
 with gr.Blocks() as demo:
     gr.Markdown("# AI-Powered Video Stabilization")
-    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.")
+    gr.Markdown(
+        "Upload a video, select a zoom factor, choose whether to apply only vertical stabilization, and optionally compress the video before processing. "
+        "The system will generate motion data using an AI model (RAFT if available) and then stabilize the video with live progress updates and alerts."
+    )
     
     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)
+            compress_checkbox = gr.Checkbox(label="Compress Video Before Processing", value=False)
+            compression_slider = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.5, label="Compression Factor (Scale)")
             process_button = gr.Button("Process Video")
         with gr.Column():
             original_video = gr.Video(label="Original Video")
@@ -254,7 +318,7 @@ with gr.Blocks() as demo:
     
     process_button.click(
         fn=process_video_ai,
-        inputs=[video_input, zoom_slider, vertical_checkbox],
+        inputs=[video_input, zoom_slider, vertical_checkbox, compress_checkbox, compression_slider],
         outputs=[original_video, stabilized_video, logs_output]
     )