#!/usr/bin/env python3 from __future__ import annotations import argparse import csv import sys from pathlib import Path from typing import Optional, Sequence import numpy as np from PIL import Image PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from data.instant_gaussian_nersemble_dataset import InstantGaussianNersembleDataset from model.xnemo_pose_preprocess import FaceNotFoundError, XNemoPosePreprocessor def main() -> int: args = parse_args() subjects = parse_list(args.subjects) cameras = parse_list(args.cameras) sequences = parse_list(args.sequences) dataset = InstantGaussianNersembleDataset( root=args.root, split=args.split, subjects=subjects, cameras=cameras, sequences=sequences, frame_stride=args.frame_stride, max_frames_per_sequence=args.max_frames_per_sequence, pose_mode="none", return_metadata=False, ) records = list(dataset.records) if args.num_shards <= 0: raise ValueError("--num-shards must be positive.") if not 0 <= args.shard_index < args.num_shards: raise ValueError("--shard-index must be in [0, num_shards).") if args.num_shards > 1: records = [record for i, record in enumerate(records) if i % args.num_shards == args.shard_index] if args.limit is not None: records = records[: args.limit] cache_root = Path(args.pose_cache_root) preview_root = Path(args.preview_root) if args.preview_root is not None else None failure_log = Path(args.failure_log) if args.failure_log is not None else None if failure_log is not None: failure_log.parent.mkdir(parents=True, exist_ok=True) print(f"records: {len(records)}") print(f"pose cache root: {cache_root}") if args.dry_run: for record in records[: min(10, len(records))]: print(record_to_cache_path(cache_root, record)) return 0 preprocessor = XNemoPosePreprocessor( xnemo_root=args.xnemo_root, output_size=args.output_size, backend=args.backend, crop=not args.no_crop, bbox_scale=args.bbox_scale, bbox_param_mode=args.bbox_param_mode, min_detection_confidence=args.min_detection_confidence, ) failures: list[tuple[str, str]] = [] counters = {"done": 0, "skipped": 0, "failed": 0} iterator = progress(records, enabled=not args.no_progress) try: for record in iterator: cache_path = record_to_cache_path(cache_root, record) if cache_path.exists() and not args.overwrite: counters["skipped"] += 1 continue try: image = dataset._load_image(record) result = preprocessor(image) except FaceNotFoundError as exc: counters["failed"] += 1 failures.append((str(record.processed_dir), f"face_not_found: {exc}")) if args.fail_fast: raise continue except Exception as exc: counters["failed"] += 1 failures.append((str(record.processed_dir), repr(exc))) if args.fail_fast: raise continue pose_tensor = result.pose_tensor.float().cpu().numpy().astype(np.float32) payload = { "pose_image": result.pose_image.astype(np.uint8), "pose_tensor": pose_tensor, "bbox_param": result.bbox_param.cpu().numpy().astype(np.float32), "bbox": result.bbox.astype(np.float32), "crop_bbox": result.crop_bbox.astype(np.float32), "landmarks": result.landmarks.astype(np.float32), } cache_path.parent.mkdir(parents=True, exist_ok=True) np.savez_compressed(cache_path, **payload) counters["done"] += 1 if preview_root is not None: preview_path = record_to_cache_path(preview_root, record).with_suffix(".png") preview_path.parent.mkdir(parents=True, exist_ok=True) Image.fromarray(result.pose_image).save(preview_path) finally: preprocessor.close() if failure_log is not None and failures: write_failures(failure_log, failures) print( "finished: " f"done={counters['done']} " f"skipped={counters['skipped']} " f"failed={counters['failed']}" ) if failures: print(f"failure log: {failure_log if failure_log is not None else '(not written)'}") return 1 if failures and args.fail_on_any_failure else 0 def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Precompute X-NeMo face-mesh RGB pose image cache for FastAvatar NeRSemble data." ) parser.add_argument("--root", default="/workspace/cpfs-data/data/nersemble_fastavatar") parser.add_argument( "--pose-cache-root", default="/workspace/cpfs-data/data/nersemble_fastavatar_xnemo_pose_cache", ) parser.add_argument("--xnemo-root", default="/workspace/cpfs-data/code/gha/x-nemo-inference") parser.add_argument("--split", default="all", choices=["all", "train", "val", "test"]) parser.add_argument("--subjects", nargs="*", default=None) parser.add_argument("--cameras", nargs="*", default=None) parser.add_argument("--sequences", nargs="*", default=None) parser.add_argument("--frame-stride", type=int, default=1) parser.add_argument("--max-frames-per-sequence", type=int, default=None) parser.add_argument("--limit", type=int, default=None) parser.add_argument("--num-shards", type=int, default=1) parser.add_argument("--shard-index", type=int, default=0) parser.add_argument("--backend", default="auto", choices=["auto", "xnemo_task", "mediapipe"]) parser.add_argument("--output-size", type=int, default=224) parser.add_argument("--bbox-scale", type=float, default=1.1) parser.add_argument("--bbox-param-mode", default="identity", choices=["identity", "computed"]) parser.add_argument("--min-detection-confidence", type=float, default=0.5) parser.add_argument("--no-crop", action="store_true") parser.add_argument("--overwrite", action="store_true") parser.add_argument("--dry-run", action="store_true") parser.add_argument("--fail-fast", action="store_true") parser.add_argument("--fail-on-any-failure", action="store_true") parser.add_argument("--failure-log", default="preprocess_xnemo_pose_failures.csv") parser.add_argument("--preview-root", default=None) parser.add_argument("--no-progress", action="store_true") return parser.parse_args() def parse_list(values: Optional[Sequence[str]]) -> Optional[list[str]]: if values is None: return None out: list[str] = [] for value in values: for item in str(value).split(","): item = item.strip() if item: out.append(item) return out or None def record_to_cache_path(root: Path, record) -> Path: return ( root / record.subject / record.camera_dir / record.sequence / f"{record.frame_index:05d}.npz" ) def write_failures(path: Path, failures: Sequence[tuple[str, str]]) -> None: path.parent.mkdir(parents=True, exist_ok=True) with open(path, "w", newline="") as f: writer = csv.writer(f) writer.writerow(["processed_dir", "error"]) writer.writerows(failures) def progress(records, enabled: bool): if not enabled: return records try: from tqdm import tqdm return tqdm(records, desc="precompute pose", dynamic_ncols=True) except Exception: return records if __name__ == "__main__": raise SystemExit(main())