Update app.py

app.py

@@ -27,10 +27,12 @@ except Exception as e:
     print("[INFO] Falling back to OpenCV Farneback optical flow.")
     raft_model = None
 
-def generate_motion_csv(video_file, output_csv=None):
+def generate_motion_csv(video_file, output_csv=None, progress=gr.Progress(), progress_offset=0.0, progress_scale=0.5):
     """
     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.
+    
+    The progress bar is updated from progress_offset to progress_offset+progress_scale.
     """
     start_time = time.time()
     if output_csv is None:
@@ -107,6 +109,9 @@ def generate_motion_csv(video_file, output_csv=None):
             
             if frame_idx % 10 == 0 or frame_idx == total_frames:
                 print(f"[INFO] Processed frame {frame_idx}/{total_frames}")
+            
+            # Update progress for this phase.
+            progress(progress_offset + (frame_idx / total_frames) * progress_scale, desc="Generating Motion CSV")
             frame_idx += 1
 
     cap.release()
@@ -141,10 +146,12 @@ 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, output_file=None):
+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):
     """
     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).
+    
+    The progress bar is updated from progress_offset to progress_offset+progress_scale.
     """
     start_time = time.time()
     print(f"[INFO] Starting stabilization using CSV: {csv_file}")
@@ -185,8 +192,7 @@ def stabilize_video_using_csv(video_file, csv_file, zoom=1.0, vertical_only=Fals
         
         dx, dy = motion_data.get(frame_idx, (0, 0))
         if vertical_only:
-            # If only vertical stabilization is desired, ignore horizontal motion.
-            dx = 0
+            dx = 0  # Ignore horizontal motion for vertical-only stabilization.
         transform = np.array([[1, 0, dx],
                               [0, 1, dy]], dtype=np.float32)
         stabilized_frame = cv2.warpAffine(frame, transform, (width, height))
@@ -194,6 +200,9 @@ def stabilize_video_using_csv(video_file, csv_file, zoom=1.0, vertical_only=Fals
         out.write(stabilized_frame)
         if frame_idx % 10 == 0 or frame_idx == total_frames:
             print(f"[INFO] Stabilized frame {frame_idx}/{total_frames}")
+        
+        # Update progress for stabilization phase.
+        progress(progress_offset + (frame_idx / total_frames) * progress_scale, desc="Stabilizing Video")
         frame_idx += 1
     
     cap.release()
@@ -202,7 +211,7 @@ 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):
+def process_video_ai(video_file, zoom, vertical_only, progress=gr.Progress(track_tqdm=True)):
     """
     Gradio interface function:
       - Generates motion data (CSV) from the input video using an AI model (RAFT if available, else Farneback).
@@ -220,8 +229,11 @@ def process_video_ai(video_file, zoom, vertical_only):
             raise ValueError("[ERROR] Please upload a video file.")
     
         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, vertical_only=vertical_only)
+        # First half: Generate motion CSV.
+        csv_file = generate_motion_csv(video_file, progress=progress, progress_offset=0.0, progress_scale=0.5)
+        # Second half: Stabilize video.
+        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)
         print("[INFO] Video processing complete.")
     logs = log_buffer.getvalue()
     return video_file, stabilized_path, logs
@@ -229,7 +241,7 @@ def process_video_ai(video_file, zoom, vertical_only):
 # 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 accordingly.")
+    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 with live progress updates.")
     
     with gr.Row():
         with gr.Column():