From e3f6def234d4f01a27881f9af28cf46a4bab787c Mon Sep 17 00:00:00 2001
From: OleehyO <leehy0357@gmail.com>
Date: Fri, 27 Dec 2024 09:50:47 +0000
Subject: [PATCH] feat: add video frame extraction tool

Add utility script to extract first frames from videos, helping users convert T2V datasets to I2V format
---
 finetune/scripts/extract_images.py | 52 ++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)
 create mode 100644 finetune/scripts/extract_images.py

diff --git a/finetune/scripts/extract_images.py b/finetune/scripts/extract_images.py
new file mode 100644
index 0000000..a5cabae
--- /dev/null
+++ b/finetune/scripts/extract_images.py
@@ -0,0 +1,52 @@
+import argparse
+import os
+from pathlib import Path
+import cv2
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--datadir", type=str, required=True, help="Root directory containing videos.txt and video subdirectory")
+    return parser.parse_args()
+
+args = parse_args()
+
+# Create data/images directory if it doesn't exist
+data_dir = Path(args.datadir)
+image_dir = data_dir / "images"
+image_dir.mkdir(exist_ok=True)
+
+# Read videos.txt
+videos_file = data_dir / "videos.txt"
+with open(videos_file, "r") as f:
+    video_paths = [line.strip() for line in f.readlines() if line.strip()]
+
+# Process each video file and collect image paths
+image_paths = []
+for video_rel_path in video_paths:
+    video_path = data_dir / video_rel_path
+    
+    # Open video
+    cap = cv2.VideoCapture(str(video_path))
+    
+    # Read first frame
+    ret, frame = cap.read()
+    if not ret:
+        print(f"Failed to read video: {video_path}")
+        continue
+        
+    # Save frame as PNG with same name as video
+    image_name = f"images/{video_path.stem}.png"
+    image_path = data_dir / image_name
+    cv2.imwrite(str(image_path), frame)
+    
+    # Release video capture
+    cap.release()
+    
+    print(f"Extracted first frame from {video_path} to {image_path}")
+    image_paths.append(image_name)
+
+# Write images.txt
+images_file = data_dir / "images.txt"
+with open(images_file, "w") as f:
+    for path in image_paths:
+        f.write(f"{path}\n")
\ No newline at end of file