File size: 3,503 Bytes
e23cf4c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | """Trích frame từ video REAL (FF++ original / Celeb-DF real) về ảnh 256x256,
theo đúng protocol của DF40: lấy mỗi frame thứ 6 (select='not(mod(n,6))'),
resize 256x256, tên file giữ số frame gốc (000.png, 006.png, 012.png, ...).
Đích: tạo phần "real" ghép với fake cdf/ff của các method chỉ-có-fake trong
DF40, cho benchmark cân bằng real/fake.
Chạy:
.venv/bin/python scripts/extract_real_frames.py \
--videos data/real-root/original_sequences/youtube/c23/videos/585.mp4 \
--out data/real-root/frames/ff++/585 \
--step 6 --size 256
# nhiều video:
find data/real-root/original_sequences/youtube/c23/videos -name '*.mp4' | \
head -50 | xargs -I{} .venv/bin/python scripts/extract_real_frames.py --videos {} --out data/real-root/frames/ff++
"""
import argparse
import os
import subprocess
import sys
FFMPEG = "ffmpeg"
FFPROBE = "ffprobe"
def frame_count(video):
"""Số frame video (dùng để biết rename đúng)."""
out = subprocess.run(
[FFPROBE, "-v", "error", "-select_streams", "v:0",
"-count_frames", "-show_entries", "stream=nb_read_frames",
"-of", "csv=p=0", video],
capture_output=True, text=True)
try:
return int(out.stdout.strip().split("\n")[0])
except (ValueError, IndexError):
return None
def extract(video, out_dir, step, size):
"""Trích frame mỗi `step` frame, resize `size`x`size`, giữ số frame gốc."""
os.makedirs(out_dir, exist_ok=True)
n_total = frame_count(video)
if n_total is None:
print(f" [skip] không đọc được frame count: {video}")
return 0
tmp = os.path.join(out_dir, ".tmp_%04d.png")
vf = f"select='not(mod(n\\,{step}))',scale={size}:{size}:flags=bicubic"
r = subprocess.run(
[FFMPEG, "-y", "-loglevel", "error", "-i", video,
"-vf", vf, "-vsync", "0", tmp],
capture_output=True, text=True)
if r.returncode != 0:
print(f" [skip] ffmpeg lỗi {video}: {r.stderr.strip()[:200]}")
return 0
# rename .tmp_XXXX.png -> <index gốc>.png
frames = sorted(f for f in os.listdir(out_dir) if f.startswith(".tmp_"))
for i, f in enumerate(frames):
src_idx = i * step # frame gốc (0, step, 2*step, ...)
os.rename(os.path.join(out_dir, f),
os.path.join(out_dir, f"{src_idx:03d}.png"))
print(f" {os.path.basename(video)}: {len(frames)} frames -> {out_dir}")
return len(frames)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--videos", required=True, help="video(s) mp4, phân cách bằng dấu phẩy hoặc lặp lại --videos")
ap.add_argument("--out", required=True, help="thư mục xuất (mỗi video 1 thư mục con)")
ap.add_argument("--step", type=int, default=6)
ap.add_argument("--size", type=int, default=256)
args = ap.parse_args()
vids = [v for v in args.videos.split(",") if v]
total = 0
for v in vids:
if not os.path.isfile(v):
print(f" [thiếu] {v}")
continue
# out/<tên video gốc>/<frame>.png
vid_dir = os.path.join(args.out, os.path.splitext(os.path.basename(v))[0])
if os.path.exists(vid_dir) and len(os.listdir(vid_dir)):
print(f" [skip đã có] {vid_dir}")
continue
total += extract(v, vid_dir, args.step, args.size)
print(f"\nTổng frame trích: {total:,}")
if __name__ == "__main__":
main()
|