optical_flow_dense.py

import cv2
import numpy as np
from pathlib import Path

# ---------- CONFIG ----------
INPUT_VIDEO = "/home/workdir/attachments/1000122934.mp4"
OUTPUT_VIDEO = "/home/workdir/artifacts/optical_flow_output.mp4"
# ----------------------------

def main():
    cap = cv2.VideoCapture(INPUT_VIDEO)
    if not cap.isOpened():
        print(f"Error opening {INPUT_VIDEO}")
        return

    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps    = cap.get(cv2.CAP_PROP_FPS) or 30.0
    total  = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    print(f"Video: {width}x{height} @ {fps:.1f} fps, {total} frames")

    # Output writer
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = cv2.VideoWriter(OUTPUT_VIDEO, fourcc, fps, (width, height))

    # First frame
    ret, prev = cap.read()
    if not ret:
        print("Could not read first frame")
        return
    prev_gray = cv2.cvtColor(prev, cv2.COLOR_BGR2GRAY)

    # HSV image for flow visualization
    hsv = np.zeros_like(prev)
    hsv[..., 1] = 255  # full saturation

    frame_idx = 0
    magnitudes = []

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # Dense optical flow (Farneback)
        flow = cv2.calcOpticalFlowFarneback(
            prev_gray, gray,
            None,
            pyr_scale=0.5,
            levels=3,
            winsize=15,
            iterations=3,
            poly_n=5,
            poly_sigma=1.2,
            flags=0
        )

        # Convert flow to polar coordinates (magnitude + angle)
        mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1])
        hsv[..., 0] = ang * 180 / np.pi / 2          # hue = direction
        hsv[..., 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)  # value = speed
        flow_rgb = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)

        # Blend original frame + flow visualization for context
        blended = cv2.addWeighted(frame, 0.4, flow_rgb, 0.6, 0)

        # Optional: draw a few sparse arrows for clarity (every 20 px)
        step = 20
        for y in range(0, height, step):
            for x in range(0, width, step):
                fx, fy = flow[y, x]
                if abs(fx) + abs(fy) > 1.5:  # only significant motion
                    cv2.arrowedLine(
                        blended,
                        (x, y),
                        (int(x + fx), int(y + fy)),
                        (0, 255, 255),
                        1,
                        tipLength=0.3
                    )

        out.write(blended)

        # Collect magnitude stats (ignore near-zero)
        valid_mag = mag[mag > 0.5]
        if len(valid_mag) > 0:
            magnitudes.append(float(np.mean(valid_mag)))

        prev_gray = gray
        frame_idx += 1

        if frame_idx % 50 == 0:
            print(f"Processed {frame_idx}/{total} frames...")

    cap.release()
    out.release()

    print(f"\nOptical flow video saved → {OUTPUT_VIDEO}")
    print(f"Total frames processed: {frame_idx}")
    if magnitudes:
        print(f"Average motion magnitude (where motion > 0.5): {np.mean(magnitudes):.2f}")
        print(f"Peak average magnitude in a frame: {np.max(magnitudes):.2f}")
        print(f"Frames with measurable motion: {len(magnitudes)}")

if __name__ == "__main__":
    main()