Update app.py

app.py

@@ -15,7 +15,6 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"[INFO] Using device: {device}")
 
 # Try to load the RAFT model from torch.hub.
-# If it fails, we fall back to OpenCV optical flow.
 try:
     print("[INFO] Attempting to load RAFT model from torch.hub...")
     raft_model = torch.hub.load("princeton-vl/RAFT", "raft_small", pretrained=True, trust_repo=True)
@@ -26,13 +25,14 @@ except Exception as e:
     print("[ERROR] Error loading RAFT model:", e)
     print("[INFO] Falling back to OpenCV Farneback optical flow.")
     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):
     """
     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.
+    Updates progress from progress_offset to progress_offset+progress_scale.
     """
     start_time = time.time()
     if output_csv is None:
@@ -42,7 +42,7 @@ def generate_motion_csv(video_file, output_csv=None, progress=gr.Progress(), pro
     
     cap = cv2.VideoCapture(video_file)
     if not cap.isOpened():
-        raise ValueError("[ERROR] Could not open video file for CSV generation.")
+        raise gr.Error("Could not open video file for CSV generation.")
     
     print(f"[INFO] Generating motion CSV for video: {video_file}")
     with open(output_csv, 'w', newline='') as csvfile:
@@ -52,7 +52,7 @@ def generate_motion_csv(video_file, output_csv=None, progress=gr.Progress(), pro
         
         ret, first_frame = cap.read()
         if not ret:
-            raise ValueError("[ERROR] Cannot read first frame from video.")
+            raise gr.Error("Cannot read first frame from video.")
         
         if raft_model is not None:
             first_frame_rgb = cv2.cvtColor(first_frame, cv2.COLOR_BGR2RGB)
@@ -86,12 +86,12 @@ 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 of the optical flow.
+            # 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 a "zoom factor": fraction of pixels moving away from the center.
+            # 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))
@@ -110,7 +110,6 @@ def generate_motion_csv(video_file, output_csv=None, progress=gr.Progress(), pro
             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
 
@@ -121,11 +120,9 @@ def generate_motion_csv(video_file, output_csv=None, progress=gr.Progress(), pro
 
 def read_motion_csv(csv_filename):
     """
-    Reads a motion CSV file (with columns: frame, mag, ang, zoom) and computes a cumulative
-    offset per frame for stabilization.
+    Reads a motion CSV file and computes cumulative offset per frame.
     
-    Returns:
-        A dictionary mapping frame numbers to (dx, dy) offsets.
+    Returns a dictionary mapping frame numbers to (dx, dy) offsets.
     """
     print(f"[INFO] Reading motion CSV: {csv_filename}")
     motion_data = {}
@@ -148,10 +145,10 @@ def read_motion_csv(csv_filename):
 
 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).
+    Stabilizes the input video using motion data from the CSV.
+    If vertical_only is True, only vertical motion is corrected.
     
-    The progress bar is updated from progress_offset to progress_offset+progress_scale.
+    Updates progress from progress_offset to progress_offset+progress_scale.
     """
     start_time = time.time()
     print(f"[INFO] Starting stabilization using CSV: {csv_file}")
@@ -159,7 +156,7 @@ def stabilize_video_using_csv(video_file, csv_file, zoom=1.0, vertical_only=Fals
     
     cap = cv2.VideoCapture(video_file)
     if not cap.isOpened():
-        raise ValueError("[ERROR] Could not open video file for stabilization.")
+        raise gr.Error("Could not open video file for stabilization.")
     
     fps = cap.get(cv2.CAP_PROP_FPS)
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
@@ -192,7 +189,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:
-            dx = 0  # Ignore horizontal motion for vertical-only stabilization.
+            dx = 0  # Only vertical stabilization.
         transform = np.array([[1, 0, dx],
                               [0, 1, dy]], dtype=np.float32)
         stabilized_frame = cv2.warpAffine(frame, transform, (width, height))
@@ -201,7 +198,6 @@ def stabilize_video_using_csv(video_file, csv_file, zoom=1.0, vertical_only=Fals
         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
     
@@ -214,26 +210,28 @@ def stabilize_video_using_csv(video_file, csv_file, zoom=1.0, vertical_only=Fals
 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).
+      - Generates motion data 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.
+      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 ValueError("[ERROR] Please upload a video file.")
+            raise gr.Error("Please upload a video file.")
     
-        print("[INFO] Starting AI-powered video processing...")
-        # 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.
+        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)
+        gr.Info("Video stabilization complete.")
         print("[INFO] Video processing complete.")
     logs = log_buffer.getvalue()
     return video_file, stabilized_path, logs
@@ -241,7 +239,7 @@ 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 with live progress updates.")
+    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 and alerts.")
     
     with gr.Row():
         with gr.Column():