Update app.py

app.py

@@ -29,20 +29,20 @@ except Exception as e:
 def process_video_ai(video_file, zoom):
     """
     Generator function for Gradio:
-      - Generates motion data (CSV) from the input video using an AI model (RAFT if available, else Farneback)
+      - Generates motion data (CSV) from the input video using an AI model (RAFT if available, else Farneback).
       - Stabilizes the video using the generated motion data.
       
     Yields:
       A tuple of (original_video, stabilized_video, logs, progress)
       During processing, original_video and stabilized_video are None.
-      The final yield returns the video file paths along with final logs and progress=100.
+      The final yield returns the video file paths with final logs and progress=100.
     """
     logs = []
     def add_log(msg):
         logs.append(msg)
         return "\n".join(logs)
     
-    # Check and extract the file path
+    # Check and extract the file path.
     if isinstance(video_file, dict):
         video_file = video_file.get("name", None)
     if video_file is None:
@@ -63,7 +63,7 @@ def process_video_ai(video_file, zoom):
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     add_log(f"[INFO] Total frames in video: {total_frames}")
     
-    # Create temporary CSV file
+    # Create temporary CSV file.
     csv_file = tempfile.NamedTemporaryFile(delete=False, suffix='.csv').name
     with open(csv_file, 'w', newline='') as csvfile:
         fieldnames = ['frame', 'mag', 'ang', 'zoom']
@@ -85,11 +85,11 @@ def process_video_ai(video_file, zoom):
             add_log("[INFO] Using Farneback optical flow for computation.")
         
         frame_idx = 1
-        # Process each frame for CSV generation
         while True:
             ret, frame = cap.read()
             if not ret:
                 break
+
             if raft_model is not None:
                 curr_frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                 curr_tensor = torch.from_numpy(curr_frame_rgb).permute(2, 0, 1).float().unsqueeze(0) / 255.0
@@ -105,11 +105,11 @@ def process_video_ai(video_file, zoom):
                                                     iterations=3, poly_n=5, poly_sigma=1.2, flags=0)
                 prev_gray = curr_gray
             
-            # Compute median magnitude and angle
+            # 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
+            # Compute zoom factor: fraction of pixels moving away from the 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))
@@ -126,7 +126,7 @@ def process_video_ai(video_file, zoom):
             })
             
             if frame_idx % 10 == 0 or frame_idx == total_frames:
-                progress_csv = (frame_idx / total_frames) * 50  # CSV phase is 0-50%
+                progress_csv = (frame_idx / total_frames) * 50  # CSV phase: 0-50%
                 add_log(f"[INFO] CSV: Processed frame {frame_idx}/{total_frames}")
                 yield (None, None, add_log(""), progress_csv)
             frame_idx += 1
@@ -138,7 +138,7 @@ def process_video_ai(video_file, zoom):
     add_log("[INFO] Starting video stabilization...")
     yield (None, None, add_log("Starting stabilization..."), 51)
     
-    # Read the CSV and compute cumulative motion data
+    # Read the CSV and compute cumulative motion data.
     motion_data = {}
     cumulative_dx = 0.0
     cumulative_dy = 0.0
@@ -157,7 +157,7 @@ def process_video_ai(video_file, zoom):
     add_log("[INFO] Motion CSV read complete.")
     yield (None, None, add_log(""), 55)
     
-    # Re-open video for stabilization
+    # Re-open video for stabilization.
     cap = cv2.VideoCapture(video_file)
     fps = cap.get(cv2.CAP_PROP_FPS)
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
@@ -182,12 +182,13 @@ def process_video_ai(video_file, zoom):
             frame = zoomed_frame[start_y:start_y+height, start_x:start_x+width]
         
         dx, dy = motion_data.get(frame_idx, (0, 0))
-        transform = np.array([[1, 0, dx], [0, 1, dy]], dtype=np.float32)
+        transform = np.array([[1, 0, dx],
+                              [0, 1, dy]], dtype=np.float32)
         stabilized_frame = cv2.warpAffine(frame, transform, (width, height))
         out.write(stabilized_frame)
         
         if frame_idx % 10 == 0 or frame_idx == total_frames:
-            progress_stab = 50 + (frame_idx / total_frames) * 50  # Stabilization phase is 50-100%
+            progress_stab = 50 + (frame_idx / total_frames) * 50  # Stabilization phase: 50-100%
             add_log(f"[INFO] Stabilization: Processed frame {frame_idx}/{total_frames}")
             yield (None, None, add_log(""), progress_stab)
         frame_idx += 1
@@ -196,7 +197,7 @@ def process_video_ai(video_file, zoom):
     add_log("[INFO] Stabilization complete.")
     yield (video_file, output_file, add_log(""), 100)
 
-# Build the Gradio UI with streaming enabled.
+# 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, else Farneback) and then stabilize the video. Logs and progress will update during processing.")
@@ -212,13 +213,22 @@ with gr.Blocks() as demo:
             logs_output = gr.Textbox(label="Logs", lines=15)
             progress_bar = gr.Slider(label="Progress", minimum=0, maximum=100, value=0, interactive=False)
     
-    demo.queue()  # enable streaming
+    demo.queue()  # enable queue for streaming
     
-    process_button.click(
-        fn=process_video_ai,
-        inputs=[video_input, zoom_slider],
-        outputs=[original_video, stabilized_video, logs_output, progress_bar],
-        stream=True  # enable streaming updates
-    )
+    # Try using stream=True. If that raises a TypeError, fall back without it.
+    try:
+        process_button.click(
+            fn=process_video_ai,
+            inputs=[video_input, zoom_slider],
+            outputs=[original_video, stabilized_video, logs_output, progress_bar],
+            stream=True
+        )
+    except TypeError as e:
+        print("[WARNING] Streaming not supported in this version of Gradio. Disabling streaming.")
+        process_button.click(
+            fn=process_video_ai,
+            inputs=[video_input, zoom_slider],
+            outputs=[original_video, stabilized_video, logs_output, progress_bar]
+        )
 
 demo.launch()