| |
| 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()) |
|
|
