Initial commit

dagster_project/assets/video_extraction.py

@@ -0,0 +1,83 @@
+"""Asset 1: Extract frames from video"""
+
+import cv2
+from pathlib import Path
+from typing import Dict
+from dagster import asset, AssetExecutionContext
+
+
+@asset(
+    io_manager_key="json_io_manager",
+    compute_kind="opencv",
+    description="Extract frames from video starting at specified second"
+)
+def extract_video_frames(context: AssetExecutionContext) -> Dict:
+    """
+    Extract frames from DJI_0017.MP4 video
+
+    Inputs:
+        - DJI_0017.MP4 (video file in root directory)
+
+    Outputs:
+        - data/frames/frame_XXXX.jpg (100 frames)
+        - data/extract_video_frames.json (metadata)
+
+    Returns:
+        Dict with:
+            - frames_dir: path to frames directory
+            - num_frames: number of extracted frames
+            - fps: video FPS
+            - start_frame: starting frame number
+    """
+    # Configuration
+    video_path = "DJI_0017.MP4"
+    start_sec = 10
+    num_frames = 100
+
+    context.log.info(f"Opening video: {video_path}")
+    cap = cv2.VideoCapture(video_path)
+
+    if not cap.isOpened():
+        raise RuntimeError(f"Could not open video: {video_path}")
+
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    start_frame = int(start_sec * fps)
+
+    context.log.info(f"Video info: {total_frames} frames, {fps} FPS")
+    context.log.info(f"Extracting {num_frames} frames starting from frame {start_frame} ({start_sec}s)")
+
+    # Create output directory with run_id
+    run_id = context.run_id
+    frames_dir = Path(f"data/{run_id}/frames")
+    frames_dir.mkdir(parents=True, exist_ok=True)
+
+    # Set starting position
+    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
+
+    # Extract frames
+    extracted = 0
+    for i in range(num_frames):
+        ret, frame = cap.read()
+        if not ret:
+            context.log.warning(f"Could not read frame {i}. Stopping extraction.")
+            break
+
+        frame_path = frames_dir / f"frame_{i:04d}.jpg"
+        cv2.imwrite(str(frame_path), frame)
+        extracted += 1
+
+        if (i + 1) % 20 == 0:
+            context.log.info(f"Extracted {i + 1}/{num_frames} frames")
+
+    cap.release()
+
+    context.log.info(f"✓ Extracted {extracted} frames to {frames_dir}")
+
+    return {
+        "frames_dir": str(frames_dir),
+        "num_frames": extracted,
+        "fps": fps,
+        "start_frame": start_frame,
+        "start_sec": start_sec
+    }