145 lines
4.6 KiB
Python
145 lines
4.6 KiB
Python
import argparse
|
|
import os
|
|
import random
|
|
import shutil
|
|
from typing import Dict, List, Sequence, Tuple
|
|
|
|
from tqdm import tqdm
|
|
|
|
from src.dataset import scan_videos, build_manifest
|
|
|
|
|
|
def sanitize(name: str) -> str:
|
|
bad = [":", "\\", "/", "*", "?", "\"", "<", ">", "|"]
|
|
for ch in bad:
|
|
name = name.replace(ch, "_")
|
|
return name
|
|
|
|
|
|
def parse_root_limits(root_limits: Sequence[str]) -> List[Tuple[str, int]]:
|
|
pairs = []
|
|
for item in root_limits:
|
|
if "=" not in item:
|
|
raise ValueError(f"Invalid --root-limit '{item}', expected PATH=NUM")
|
|
path, num = item.rsplit("=", 1)
|
|
pairs.append((path, int(num)))
|
|
return pairs
|
|
|
|
|
|
def choose_files(files: List[str], limit: int, seed: int) -> List[str]:
|
|
if limit <= 0 or limit >= len(files):
|
|
return files
|
|
rng = random.Random(seed)
|
|
files = files[:]
|
|
rng.shuffle(files)
|
|
return files[:limit]
|
|
|
|
|
|
def copy_mode(src: str, dst: str, mode: str) -> None:
|
|
if mode == "copy":
|
|
shutil.copy2(src, dst)
|
|
return
|
|
if mode == "hardlink":
|
|
os.link(src, dst)
|
|
return
|
|
if mode == "symlink":
|
|
os.symlink(src, dst)
|
|
return
|
|
raise ValueError(f"Unknown mode: {mode}")
|
|
|
|
|
|
def prepare_dataset(
|
|
out_dir: str,
|
|
root_limits: List[Tuple[str, int]],
|
|
extensions: Sequence[str],
|
|
seed: int,
|
|
mode: str,
|
|
skip_existing: bool,
|
|
max_total: int,
|
|
dry_run: bool,
|
|
) -> List[str]:
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
selected_all: List[str] = []
|
|
total_count = 0
|
|
for idx, (root, limit) in enumerate(root_limits):
|
|
root = os.path.normpath(root)
|
|
files = scan_videos([root], extensions)
|
|
chosen = choose_files(files, limit, seed + idx)
|
|
if max_total > 0:
|
|
remaining = max_total - total_count
|
|
if remaining <= 0:
|
|
break
|
|
chosen = chosen[:remaining]
|
|
root_name = f"root_{idx+1:02d}_{sanitize(os.path.basename(root))}"
|
|
for src in chosen:
|
|
rel = os.path.relpath(src, root)
|
|
dst = os.path.join(out_dir, root_name, rel)
|
|
selected_all.append((src, dst))
|
|
total_count += len(chosen)
|
|
to_copy = selected_all
|
|
if dry_run:
|
|
return [dst for _, dst in to_copy]
|
|
|
|
pbar = tqdm(to_copy, desc="copy", total=len(to_copy))
|
|
copied = []
|
|
for src, dst in pbar:
|
|
if skip_existing and os.path.exists(dst):
|
|
copied.append(dst)
|
|
continue
|
|
os.makedirs(os.path.dirname(dst), exist_ok=True)
|
|
try:
|
|
copy_mode(src, dst, mode)
|
|
copied.append(dst)
|
|
except Exception as exc:
|
|
pbar.set_postfix({"err": str(exc)[:80]})
|
|
continue
|
|
return copied
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--out", type=str, required=True)
|
|
parser.add_argument("--root-limit", type=str, nargs="*", default=[])
|
|
parser.add_argument("--root", type=str, nargs="*", default=[])
|
|
parser.add_argument("--limit", type=int, nargs="*", default=[])
|
|
parser.add_argument("--extensions", type=str, nargs="*", default=None)
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
parser.add_argument("--mode", type=str, default="copy", choices=["copy", "hardlink", "symlink"])
|
|
parser.add_argument("--skip-existing", action="store_true")
|
|
parser.add_argument("--max-total", type=int, default=0)
|
|
parser.add_argument("--dry-run", action="store_true")
|
|
parser.add_argument("--write-manifest", type=str, default=None)
|
|
parser.add_argument("--manifest-backend", type=str, default="opencv")
|
|
args = parser.parse_args()
|
|
|
|
if args.root_limit:
|
|
root_limits = parse_root_limits(args.root_limit)
|
|
else:
|
|
if not args.root:
|
|
raise ValueError("Specify --root-limit or --root")
|
|
if args.limit and len(args.limit) != len(args.root):
|
|
raise ValueError("--limit count must match --root count")
|
|
limits = args.limit or [0 for _ in args.root]
|
|
root_limits = list(zip(args.root, limits))
|
|
|
|
extensions = args.extensions or [".mp4", ".mkv", ".avi", ".webm", ".flv", ".mov", ".mpg", ".mpeg"]
|
|
copied = prepare_dataset(
|
|
out_dir=args.out,
|
|
root_limits=root_limits,
|
|
extensions=extensions,
|
|
seed=args.seed,
|
|
mode=args.mode,
|
|
skip_existing=args.skip_existing,
|
|
max_total=int(args.max_total),
|
|
dry_run=args.dry_run,
|
|
)
|
|
|
|
print(f"Prepared {len(copied)} files in {args.out}")
|
|
if args.write_manifest and not args.dry_run:
|
|
build_manifest([args.out], args.write_manifest, extensions, backend=args.manifest_backend)
|
|
print(f"Wrote manifest to {args.write_manifest}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|