94 lines
3.2 KiB
Python
94 lines
3.2 KiB
Python
"""Build the canonical fast-training format for AVA, WASD, and AVCC."""
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
from preprocess.asd.adapters import clip_directory_records, manifest_records
|
|
from preprocess.asd.writer import write_records
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("dataset", choices=["ava", "wasd", "msdwild", "avcc"])
|
|
parser.add_argument(
|
|
"--loader", type=Path, required=True, help="source TSV/manifest"
|
|
)
|
|
parser.add_argument(
|
|
"--output",
|
|
type=Path,
|
|
help="defaults to <dataset root>/packed/<split>, beside the source media",
|
|
)
|
|
parser.add_argument("--visual-root", type=Path)
|
|
parser.add_argument("--audio-root", type=Path)
|
|
parser.add_argument("--root", type=Path, help="base for AVCC manifest paths")
|
|
parser.add_argument("--split", default="train")
|
|
parser.add_argument("--audio-sample-rate", type=int, default=16_000)
|
|
parser.add_argument("--max-shard-gb", type=float, default=1.0)
|
|
parser.add_argument(
|
|
"--skip-missing",
|
|
action="store_true",
|
|
help="development subsets only; production preprocessing should fail on missing files",
|
|
)
|
|
return parser
|
|
|
|
|
|
def require_roots(args: argparse.Namespace) -> None:
|
|
if args.visual_root is None or args.audio_root is None:
|
|
raise ValueError(f"{args.dataset} requires --visual-root and --audio-root")
|
|
|
|
|
|
def default_output(args: argparse.Namespace) -> Path:
|
|
"""Keep each corpus's pack beside that corpus.
|
|
|
|
`<root>/clips_videos/<split>` becomes `<root>/packed/<split>`, so a dataset
|
|
stays self-contained and can be copied to fast local storage on its own.
|
|
"""
|
|
if args.visual_root is not None:
|
|
return args.visual_root.parent / "packed" / args.split
|
|
if args.root is not None:
|
|
return args.root / "packed" / args.split
|
|
raise ValueError("--output is required when no dataset root is given")
|
|
|
|
|
|
def main() -> None:
|
|
args = build_parser().parse_args()
|
|
if args.output is None:
|
|
args.output = default_output(args)
|
|
print(f"packing into {args.output}")
|
|
if args.output.exists() and any(args.output.iterdir()):
|
|
raise FileExistsError(f"output directory is not empty: {args.output}")
|
|
|
|
if args.dataset in ("ava", "wasd"):
|
|
require_roots(args)
|
|
records = clip_directory_records(
|
|
args.loader,
|
|
args.visual_root,
|
|
args.audio_root,
|
|
args.split,
|
|
dataset=args.dataset,
|
|
skip_missing=args.skip_missing,
|
|
)
|
|
elif args.dataset == "msdwild":
|
|
raise SystemExit(
|
|
"MSDWild is packed from raw media by preprocess/asd/prep_msdwild.py, "
|
|
"which needs all.rttm for labels; there is no loader TSV to read here"
|
|
)
|
|
else:
|
|
records = manifest_records(
|
|
args.loader,
|
|
dataset="avcc",
|
|
root=args.root,
|
|
default_npy_audio_sample_rate=args.audio_sample_rate,
|
|
)
|
|
|
|
count = write_records(
|
|
records,
|
|
args.output,
|
|
max_shard_bytes=round(args.max_shard_gb * 1_000_000_000),
|
|
)
|
|
print(f"Wrote {count} {args.dataset} samples to {args.output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|