Speed up

app.py

@@ -1,21 +1,18 @@
 import cv2
 import numpy as np
-import csv
-import math
 import torch
 import tempfile
-import os
 import gradio as gr
 import time
 import io
 from contextlib import redirect_stdout
-import concurrent.futures
 
-# Set up device for torch
 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 device.type == "cuda":
+    torch.backends.cudnn.benchmark = True
+
 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)
@@ -28,248 +25,203 @@ except Exception as e:
     raft_model = None
     gr.Warning("Falling back to OpenCV Farneback optical flow.")
 
-def compress_video(video_file, target_width, target_height, progress=gr.Progress(), progress_offset=0.0, progress_scale=0.2, output_file=None):
-    """
-    Compresses the video by resizing each frame to the specified target resolution.
-    The new resolution is exactly (target_width, target_height).
-    Updates progress from progress_offset to progress_offset+progress_scale.
-    """
-    start_time = time.time()
+def _resize(frame, w, h):
+    if frame.shape[1] == w and frame.shape[0] == h:
+        return frame
+    return cv2.resize(frame, (w, h), interpolation=cv2.INTER_AREA if (w < frame.shape[1] or h < frame.shape[0]) else cv2.INTER_LINEAR)
+
+def _frame_to_raft_tensor_bgr(frame_bgr):
+    frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+    t = torch.from_numpy(frame_rgb).permute(2, 0, 1).contiguous().float().unsqueeze(0).div_(255.0)
+    return t.to(device, non_blocking=(device.type == "cuda"))
+
+def compute_offsets(
+    video_file,
+    out_w,
+    out_h,
+    motion_scale=0.5,
+    raft_iters=12,
+    progress=gr.Progress(),
+    progress_offset=0.0,
+    progress_scale=0.55,
+):
     cap = cv2.VideoCapture(video_file)
     if not cap.isOpened():
-        raise gr.Error("Could not open video file for compression.")
-    
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    new_width = int(target_width)
-    new_height = int(target_height)
-    
-    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}")
+        raise gr.Error("Could not open video file for motion estimation.")
+
+    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 0
+
+    mw = max(64, int(out_w * float(motion_scale)))
+    mh = max(64, int(out_h * float(motion_scale)))
+    sx = float(out_w) / float(mw)
+    sy = float(out_h) / float(mh)
+
+    ret, prev = cap.read()
+    if not ret:
+        cap.release()
+        raise gr.Error("Cannot read first frame from video.")
+
+    prev_out = _resize(prev, out_w, out_h)
+    prev_small = _resize(prev_out, mw, mh)
+
+    use_raft = raft_model is not None
+    use_amp = device.type == "cuda"
+
+    if use_raft:
+        prev_t = _frame_to_raft_tensor_bgr(prev_small)
+    else:
+        prev_g = cv2.cvtColor(prev_small, cv2.COLOR_BGR2GRAY)
+
+    offsets = [(0.0, 0.0)]
+    cum_dx = 0.0
+    cum_dy = 0.0
+
+    idx = 1
     while True:
         ret, frame = cap.read()
         if not ret:
             break
-        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
+        frame_out = _resize(frame, out_w, out_h)
+        curr_small = _resize(frame_out, mw, mh)
 
-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()
-    if output_csv is None:
-        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.csv')
-        output_csv = temp_file.name
-        temp_file.close()
-    
-    cap = cv2.VideoCapture(video_file)
-    if not cap.isOpened():
-        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:
-        fieldnames = ['frame', 'mag', 'ang', 'zoom']
-        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
-        writer.writeheader()
-        
-        ret, first_frame = cap.read()
-        if not ret:
-            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)
-            prev_tensor = torch.from_numpy(first_frame_rgb).permute(2, 0, 1).float().unsqueeze(0) / 255.0
-            prev_tensor = prev_tensor.to(device)
-            print("[INFO] Using RAFT model for optical flow computation.")
+        if use_raft:
+            curr_t = _frame_to_raft_tensor_bgr(curr_small)
+            with torch.no_grad():
+                if use_amp:
+                    with torch.cuda.amp.autocast(True):
+                        _, flow_up = raft_model(prev_t, curr_t, iters=int(raft_iters), test_mode=True)
+                else:
+                    _, flow_up = raft_model(prev_t, curr_t, iters=int(raft_iters), test_mode=True)
+
+            flow = flow_up[0]
+            dx = float(flow[0].median().item())
+            dy = float(flow[1].median().item())
+            prev_t = curr_t
         else:
-            prev_gray = cv2.cvtColor(first_frame, cv2.COLOR_BGR2GRAY)
-            print("[INFO] Using OpenCV Farneback optical flow for computation.")
-        
-        frame_idx = 1
-        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        print(f"[INFO] Total frames to process: {total_frames}")
-        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
-                curr_tensor = curr_tensor.to(device)
-                with torch.no_grad():
-                    flow_low, flow_up = raft_model(prev_tensor, curr_tensor, iters=20, test_mode=True)
-                flow = flow_up[0].permute(1, 2, 0).cpu().numpy()
-                prev_tensor = curr_tensor.clone()
-            else:
-                curr_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-                flow = cv2.calcOpticalFlowFarneback(prev_gray, curr_gray, None,
-                                                    pyr_scale=0.5, levels=3, winsize=15,
-                                                    iterations=3, poly_n=5, poly_sigma=1.2, flags=0)
-                prev_gray = curr_gray
-
-            mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1], angleInDegrees=True)
-            median_mag = np.median(mag)
-            median_ang = np.median(ang)
-            
-            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))
-            x_offset = x_coords - center_x
-            y_offset = y_coords - center_y
-            dot = flow[..., 0] * x_offset + flow[..., 1] * y_offset
-            zoom_factor = np.count_nonzero(dot > 0) / (w * h)
-            
-            writer.writerow({
-                'frame': frame_idx,
-                'mag': median_mag,
-                'ang': median_ang,
-                'zoom': zoom_factor
-            })
-            
-            if frame_idx % 10 == 0 or frame_idx == total_frames:
-                print(f"[INFO] Processed frame {frame_idx}/{total_frames}")
-            
-            progress(progress_offset + (frame_idx / total_frames) * progress_scale, desc="Generating Motion CSV")
-            frame_idx += 1
+            curr_g = cv2.cvtColor(curr_small, cv2.COLOR_BGR2GRAY)
+            flow = cv2.calcOpticalFlowFarneback(
+                prev_g,
+                curr_g,
+                None,
+                pyr_scale=0.5,
+                levels=3,
+                winsize=15,
+                iterations=3,
+                poly_n=5,
+                poly_sigma=1.2,
+                flags=0,
+            )
+            dx = float(np.median(flow[..., 0]))
+            dy = float(np.median(flow[..., 1]))
+            prev_g = curr_g
+
+        dx *= sx
+        dy *= sy
+
+        cum_dx += dx
+        cum_dy += dy
+        offsets.append((-cum_dx, -cum_dy))
+
+        if total > 0 and (idx % 5 == 0 or idx == total - 1):
+            progress(progress_offset + (idx / max(1, total - 1)) * progress_scale, desc="Estimating Motion")
+
+        idx += 1
 
     cap.release()
-    elapsed = time.time() - start_time
-    print(f"[INFO] Motion CSV generated: {output_csv} in {elapsed:.2f} seconds")
-    return output_csv
-
-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}")
-    motion_data = {}
-    cumulative_dx = 0.0
-    cumulative_dy = 0.0
-    with open(csv_filename, 'r') as csvfile:
-        reader = csv.DictReader(csvfile)
-        for row in reader:
-            frame_num = int(row['frame'])
-            mag = float(row['mag'])
-            ang = float(row['ang'])
-            rad = math.radians(ang)
-            dx = mag * math.cos(rad)
-            dy = mag * math.sin(rad)
-            cumulative_dx += dx
-            cumulative_dy += dy
-            motion_data[frame_num] = (-cumulative_dx, -cumulative_dy)
-    print("[INFO] Completed reading motion CSV.")
-    return motion_data
-
-def process_frame(frame, dx, dy, zoom, vertical_only, width, height):
-    """
-    Processes a single frame for stabilization.
-    Applies translation by dx, dy (if not vertical_only), and scaling by zoom.
-    Uses cv2.BORDER_REPLICATE to avoid black borders.
-    """
-    if vertical_only:
-        dx = 0
-    # Create transformation matrix: translation + scaling
-    center = (width / 2, height / 2)
-    M = cv2.getRotationMatrix2D(center, 0, zoom)
-    M[0, 2] += dx
-    M[1, 2] += dy
-    stabilized_frame = cv2.warpAffine(frame, M, (width, height), borderMode=cv2.BORDER_REPLICATE)
-    return stabilized_frame
-
-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 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.
-    Uses cv2.BORDER_REPLICATE to fill any gaps, preventing black borders.
-    """
-    start_time = time.time()
-    print(f"[INFO] Starting stabilization using CSV: {csv_file}")
-    motion_data = read_motion_csv(csv_file)
-    
+    return offsets
+
+def compute_auto_zoom(offsets, width, height):
+    dxs = [o[0] for o in offsets] or [0.0]
+    dys = [o[1] for o in offsets] or [0.0]
+
+    left = max(0.0, -min(dxs))
+    right = max(0.0, max(dxs))
+    top = max(0.0, -min(dys))
+    bottom = max(0.0, max(dys))
+
+    safe_w = float(width) - (left + right)
+    safe_h = float(height) - (top + bottom)
+
+    zx = (float(width) / safe_w) if safe_w > 1.0 else 1.0
+    zy = (float(height) / safe_h) if safe_h > 1.0 else 1.0
+    return max(1.0, zx, zy)
+
+def stabilize_stream(
+    video_file,
+    offsets,
+    zoom=1.0,
+    vertical_only=False,
+    out_w=None,
+    out_h=None,
+    progress=gr.Progress(),
+    progress_offset=0.55,
+    progress_scale=0.45,
+    output_file=None,
+):
     cap = cv2.VideoCapture(video_file)
     if not cap.isOpened():
         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))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    print(f"[INFO] Video properties - FPS: {fps}, Width: {width}, Height: {height}")
-    
+    in_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    in_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    if out_w is None:
+        out_w = in_w
+    if out_h is None:
+        out_h = in_h
+
     if output_file is None:
-        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
         output_file = temp_file.name
         temp_file.close()
-    
-    frames = []
-    while True:
+
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    out = cv2.VideoWriter(output_file, fourcc, fps, (int(out_w), int(out_h)))
+
+    center = (float(out_w) / 2.0, float(out_h) / 2.0)
+    base = cv2.getRotationMatrix2D(center, 0.0, float(zoom))
+
+    total = len(offsets)
+    i = 0
+    while i < total:
         ret, frame = cap.read()
         if not ret:
             break
-        frames.append(frame)
+
+        frame_out = _resize(frame, int(out_w), int(out_h))
+
+        dx, dy = offsets[i]
+        if vertical_only:
+            dx = 0.0
+
+        M = base.copy()
+        M[0, 2] += float(dx)
+        M[1, 2] += float(dy)
+
+        stabilized = cv2.warpAffine(frame_out, M, (int(out_w), int(out_h)), borderMode=cv2.BORDER_REPLICATE)
+        out.write(stabilized)
+
+        if total > 0 and (i % 5 == 0 or i == total - 1):
+            progress(progress_offset + (i / max(1, total - 1)) * progress_scale, desc="Stabilizing Video")
+
+        i += 1
+
     cap.release()
-    total_frames = len(frames)
-    print(f"[INFO] Total frames to stabilize: {total_frames}")
-    # Use threads to avoid multiprocessing pickling issues on Windows/Gradio.
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        dx_list = [motion_data.get(i+1, (0, 0))[0] for i in range(len(frames))]
-        dy_list = [motion_data.get(i+1, (0, 0))[1] for i in range(len(frames))]
-        zoom_list = [zoom] * len(frames)
-        vertical_list = [vertical_only] * len(frames)
-        width_list = [width] * len(frames)
-        height_list = [height] * len(frames)
-        stabilized_frames = list(executor.map(process_frame, frames, dx_list, dy_list, zoom_list, vertical_list, width_list, height_list))
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_file, fourcc, fps, (width, height))
-    for idx, stabilized_frame in enumerate(stabilized_frames):
-        out.write(stabilized_frame)
-        frame_idx = idx + 1
-        if frame_idx % 10 == 0 or frame_idx == total_frames:
-            print(f"[INFO] Stabilized frame {frame_idx}/{total_frames}")
-        
-        progress(progress_offset + (frame_idx / total_frames) * progress_scale, desc="Stabilizing Video")
     out.release()
-    elapsed = time.time() - start_time
-    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, compress_mode, target_width, target_height, auto_zoom, progress=gr.Progress(track_tqdm=True)):
-    """
-    Gradio interface function:
-      - Optionally compresses the video if compress_mode is True, resizing it to the chosen resolution.
-      - Generates motion data from the (possibly compressed) video.
-      - If auto_zoom is enabled, computes the optimal zoom level based on the maximum cumulative displacements.
-      - 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)
-    """
+def process_video_ai(
+    video_file,
+    zoom,
+    vertical_only,
+    compress_mode,
+    target_width,
+    target_height,
+    auto_zoom,
+    progress=gr.Progress(track_tqdm=True),
+):
     gr.Info("Starting AI-powered video processing...")
     log_buffer = io.StringIO()
     with redirect_stdout(log_buffer):
@@ -277,82 +229,32 @@ def process_video_ai(video_file, zoom, vertical_only, compress_mode, target_widt
             video_file = video_file.get("name", None)
         if video_file is None:
             raise gr.Error("Please upload a video file.")
-        
-        # If compression is enabled, compress the video first.
+
+        cap = cv2.VideoCapture(video_file)
+        if not cap.isOpened():
+            raise gr.Error("Could not open uploaded video.")
+        in_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        in_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        cap.release()
+
         if compress_mode:
-            gr.Info("Compressing video before processing...")
-            video_file = compress_video(video_file, target_width, target_height, progress=progress, progress_offset=0.0, progress_scale=0.2)
-            gr.Info("Video compression complete.")
-            motion_offset = 0.2
-            motion_scale = 0.4
-            stabilization_offset = 0.6
-            stabilization_scale = 0.4
+            out_w = int(target_width)
+            out_h = int(target_height)
         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.")
-        
-        # Auto Zoom Mode: compute the optimal zoom factor to avoid black borders.
+            out_w = in_w
+            out_h = in_h
+
+        offsets = compute_offsets(
+            video_file,
+            out_w,
+            out_h,
+            motion_scale=0.5,
+            raft_iters=12,
+            progress=progress,
+            progress_offset=0.0,
+            progress_scale=0.55,
+        )
+        gr.Info("Motion estimated successfully.")
+
         if auto_zoom:
-            gr.Info("Auto Zoom Mode enabled. Computing optimal zoom factor...")
-            motion_data = read_motion_csv(csv_file)
-            # Compute separate left/right and top/bottom displacements.
-            left_disp = abs(min(v[0] for v in motion_data.values()))
-            right_disp = max(v[0] for v in motion_data.values())
-            top_disp = abs(min(v[1] for v in motion_data.values()))
-            bottom_disp = max(v[1] for v in motion_data.values())
-            cap = cv2.VideoCapture(video_file)
-            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-            cap.release()
-            safe_width = width - (left_disp + right_disp)
-            safe_height = height - (top_disp + bottom_disp)
-            zoom_x = width / safe_width if safe_width > 0 else 1.0
-            zoom_y = height / safe_height if safe_height > 0 else 1.0
-            auto_zoom_factor = max(1.0, zoom_x, zoom_y)
-            gr.Info(f"Auto zoom factor computed: {auto_zoom_factor:.2f}")
-            zoom = auto_zoom_factor
-
-        stabilized_path = stabilize_video_using_csv(video_file, csv_file, zoom=zoom, vertical_only=vertical_only,
-                                                      progress=progress, progress_offset=stabilization_offset, progress_scale=stabilization_scale)
-        gr.Info("Video stabilization complete.")
-        print("[INFO] Video processing complete.")
-    logs = log_buffer.getvalue()
-    return video_file, stabilized_path, logs
-
-# Build the Gradio UI.
-with gr.Blocks() as demo:
-    gr.Markdown("# AI-Powered Video Stabilization")
-    gr.Markdown(
-        "Upload a video, select a zoom factor (or use Auto Zoom Mode), choose whether to apply only vertical stabilization, and optionally compress the video before processing. "
-        "If compressing, specify the target resolution (width and height) for the compressed video. "
-        "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=3.0, step=0.1, value=1.0, label="Zoom Factor (ignored if Auto Zoom enabled)")
-            auto_zoom_checkbox = gr.Checkbox(label="Auto Zoom Mode", value=False)
-            vertical_checkbox = gr.Checkbox(label="Vertical Stabilization Only", value=False)
-            compress_checkbox = gr.Checkbox(label="Compress Video Before Processing", value=False)
-            target_width = gr.Number(label="Target Width (px)", value=640)
-            target_height = gr.Number(label="Target Height (px)", value=360)
-            process_button = gr.Button("Process Video")
-        with gr.Column():
-            original_video = gr.Video(label="Original Video")
-            stabilized_video = gr.Video(label="Stabilized Video")
-            logs_output = gr.Textbox(label="Logs", lines=10)
-    
-    process_button.click(
-        fn=process_video_ai,
-        inputs=[video_input, zoom_slider, vertical_checkbox, compress_checkbox, target_width, target_height, auto_zoom_checkbox],
-        outputs=[original_video, stabilized_video, logs_output]
-    )
-
-if __name__ == "__main__":
-    demo.launch()
+            z = compute_auto_zoom(offsets, out