Create app.py

app.py

@@ -0,0 +1,136 @@
+import cv2
+import numpy as np
+import csv
+import math
+import tempfile
+import os
+import gradio as gr
+
+def read_motion_csv(csv_filename):
+    """
+    Reads a CSV file with columns: frame, mag, ang, zoom.
+    For each row, computes a displacement from mag and ang and 
+    accumulates these to build a per-frame cumulative offset.
+    
+    Returns:
+        A dictionary mapping frame numbers to (dx, dy) offsets.
+        The offsets are the negative cumulative displacements,
+        which can be used to counteract the camera motion.
+    """
+    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'])
+            # Convert angle (in degrees) to radians
+            rad = math.radians(ang)
+            # Compute displacement vector from magnitude and angle
+            dx = mag * math.cos(rad)
+            dy = mag * math.sin(rad)
+            # Accumulate the displacement over frames
+            cumulative_dx += dx
+            cumulative_dy += dy
+            # Store the negative cumulative offset to counteract motion
+            motion_data[frame_num] = (-cumulative_dx, -cumulative_dy)
+    return motion_data
+
+def stabilize_video_using_csv(video_file, csv_file, zoom=1.0, output_file=None):
+    """
+    Stabilizes the input video using motion data from the CSV file.
+    
+    Args:
+        video_file (str): Path to the input video.
+        csv_file (str): Path to the CSV file generated by the detection code.
+        zoom (float): Optional zoom factor to apply before stabilization (default: 1.0).
+        output_file (str): Path for the output stabilized video. If None, a temporary file is created.
+        
+    Returns:
+        output_file (str): The path to the stabilized video file.
+    """
+    # Read motion data from CSV
+    motion_data = read_motion_csv(csv_file)
+    
+    cap = cv2.VideoCapture(video_file)
+    if not cap.isOpened():
+        raise ValueError("Could not open video file.")
+    
+    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))
+    
+    # Create a temporary file for output if not provided
+    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, (width, height))
+    
+    frame_num = 1
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        
+        # Optionally apply zoom (resize and center-crop)
+        if zoom != 1.0:
+            zoomed_frame = cv2.resize(frame, None, fx=zoom, fy=zoom, interpolation=cv2.INTER_LINEAR)
+            zoomed_h, zoomed_w = zoomed_frame.shape[:2]
+            start_x = max((zoomed_w - width) // 2, 0)
+            start_y = max((zoomed_h - height) // 2, 0)
+            frame = zoomed_frame[start_y:start_y+height, start_x:start_x+width]
+        
+        # Retrieve stabilization offset from CSV data (if available)
+        dx, dy = motion_data.get(frame_num, (0, 0))
+        
+        # Apply an affine transformation to counteract the motion
+        transform = np.array([[1, 0, dx],
+                              [0, 1, dy]], dtype=np.float32)
+        stabilized_frame = cv2.warpAffine(frame, transform, (width, height))
+        
+        out.write(stabilized_frame)
+        frame_num += 1
+    
+    cap.release()
+    out.release()
+    return output_file
+
+def process_video(video_file, csv_file, zoom):
+    """
+    Gradio interface function to stabilize a video.
+    Accepts an input video file, a motion CSV file, and a zoom factor.
+    Returns the original video and the stabilized video.
+    """
+    # Convert input file objects to file paths if needed
+    if isinstance(video_file, dict):
+        video_file = video_file["name"]
+    if isinstance(csv_file, dict):
+        csv_file = csv_file["name"]
+    
+    stabilized_path = stabilize_video_using_csv(video_file, csv_file, zoom=zoom)
+    return video_file, stabilized_path
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Video Stabilization with Motion Data")
+    with gr.Row():
+        with gr.Column():
+            video_input = gr.Video(label="Input Video")
+            csv_input = gr.File(label="Motion CSV File (e.g., video.flow.csv)", file_count="single")
+            zoom_slider = gr.Slider(minimum=1.0, maximum=2.0, step=0.1, value=1.0, label="Zoom Factor")
+            process_button = gr.Button("Stabilize Video")
+        with gr.Column():
+            original_video = gr.Video(label="Original Video")
+            stabilized_video = gr.Video(label="Stabilized Video")
+    
+    process_button.click(
+        fn=process_video, 
+        inputs=[video_input, csv_input, zoom_slider], 
+        outputs=[original_video, stabilized_video]
+    )
+
+demo.launch()