| |
| """Build a self-contained dataset_info.json from 3Deditverse_data. |
| |
| The generated info is compatible with the existing 3DEditFormer editing |
| loader, while the sample set is decided only by files present in the clean |
| encoded/condition-image tree. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| from pathlib import Path |
| from typing import Dict, Iterable, Optional, Tuple |
|
|
| import numpy as np |
|
|
|
|
| IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp"} |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument( |
| "--clean-root", |
| type=Path, |
| default=Path("/mnt/zsn/zsn_workspace/3DEditFormer/3DEditVerse/3Deditverse_data"), |
| help="Path to 3Deditverse_data.", |
| ) |
| parser.add_argument( |
| "--output", |
| type=Path, |
| default=None, |
| help="Output dataset_info.json path. Defaults to <clean-root>/dataset_info.json.", |
| ) |
| parser.add_argument( |
| "--indent", |
| type=int, |
| default=2, |
| help="JSON indentation. Use 0 for compact JSON.", |
| ) |
| parser.add_argument( |
| "--voxel-count-mode", |
| choices=("constant", "slat"), |
| default="constant", |
| help=( |
| "How to fill ori_voxel_num/edit_voxel_num. " |
| "constant is fast and sufficient for non-sparse SS training; " |
| "slat opens every slat_latent.npz and can be slow." |
| ), |
| ) |
| return parser.parse_args() |
|
|
|
|
| def iter_dir_names(path: Path) -> Iterable[str]: |
| if not path.is_dir(): |
| return [] |
| with os.scandir(path) as entries: |
| return sorted(entry.name for entry in entries if entry.is_dir(follow_symlinks=False)) |
|
|
|
|
| def first_existing(paths: Iterable[Path]) -> Optional[Path]: |
| for path in paths: |
| if path.is_file(): |
| return path |
| return None |
|
|
|
|
| def first_extra_image(path: Path, exclude: set[str]) -> Optional[Path]: |
| if not path.is_dir(): |
| return None |
| with os.scandir(path) as entries: |
| candidates = sorted( |
| Path(entry.path) |
| for entry in entries |
| if entry.is_file(follow_symlinks=False) |
| and Path(entry.name).suffix.lower() in IMAGE_EXTS |
| and entry.name not in exclude |
| ) |
| return candidates[0] if candidates else None |
|
|
|
|
| def clean_condition_pair(clean_root: Path, branch: str, key: str) -> Optional[Tuple[Path, Path]]: |
| cond_dir = clean_root / "condition-image" / branch / key |
| ori = first_existing([cond_dir / "ori_image.png"]) |
| if branch == "alpaca": |
| edit = first_existing([cond_dir / "after_edited_Flux.png"]) |
| else: |
| edit = first_existing([cond_dir / "edited_0.png"]) |
| if edit is None: |
| edit = first_extra_image(cond_dir, {"ori_image.png"}) |
| if ori is None or edit is None: |
| return None |
| return ori, edit |
|
|
|
|
| def latent_pair(encoded_dir: Path, key: str, name: str) -> Optional[Tuple[Path, Path]]: |
| ori = encoded_dir / key / "ori" / name |
| edit = encoded_dir / key / "edit" / name |
| if ori.is_file() and edit.is_file(): |
| return ori, edit |
| return None |
|
|
|
|
| def voxel_count(slat_path: Path) -> int: |
| with np.load(slat_path) as data: |
| if "coords" in data: |
| return int(data["coords"].shape[0]) |
| if "slat_coords" in data: |
| return int(data["slat_coords"].shape[0]) |
| raise KeyError(f"Expected coords or slat_coords in {slat_path}") |
|
|
|
|
| def rel(path: Path, clean_root: Path) -> str: |
| return path.relative_to(clean_root).as_posix() |
|
|
|
|
| def build_branch(clean_root: Path, branch: str, encoded_name: str, voxel_count_mode: str) -> Tuple[Dict[str, dict], Dict[str, int]]: |
| encoded_dir = clean_root / "encoded_ouput" / encoded_name |
| condition_dir = clean_root / "condition-image" / branch |
| encoded_keys = set(iter_dir_names(encoded_dir)) |
| condition_keys = set(iter_dir_names(condition_dir)) |
|
|
| records: Dict[str, dict] = {} |
| stats = { |
| "encoded_keys": len(encoded_keys), |
| "condition_keys": len(condition_keys), |
| "missing_condition": 0, |
| "missing_ss_latent": 0, |
| "missing_slat_latent": 0, |
| "voxel_read_errors": 0, |
| } |
|
|
| for key in sorted(encoded_keys & condition_keys): |
| cond_pair = clean_condition_pair(clean_root, branch, key) |
| if cond_pair is None: |
| stats["missing_condition"] += 1 |
| continue |
|
|
| ss_pair = latent_pair(encoded_dir, key, "ss_latent.npz") |
| if ss_pair is None: |
| stats["missing_ss_latent"] += 1 |
| continue |
|
|
| slat_pair = latent_pair(encoded_dir, key, "slat_latent.npz") |
| if slat_pair is None: |
| stats["missing_slat_latent"] += 1 |
| continue |
|
|
| if voxel_count_mode == "slat": |
| try: |
| ori_voxel_num = voxel_count(slat_pair[0]) |
| edit_voxel_num = voxel_count(slat_pair[1]) |
| except Exception as exc: |
| stats["voxel_read_errors"] += 1 |
| print(f"[WARN] skip {branch}/{key}: failed to read voxel count: {exc}") |
| continue |
| else: |
| ori_voxel_num = 1 |
| edit_voxel_num = 1 |
|
|
| ori_img, edit_img = cond_pair |
| records[key] = { |
| "ori_ss_latents_path": rel(ss_pair[0], clean_root), |
| "edit_ss_latents_path": rel(ss_pair[1], clean_root), |
| "ori_latents_path": rel(slat_pair[0], clean_root), |
| "edit_latents_path": rel(slat_pair[1], clean_root), |
| "ori_img_path": rel(ori_img, clean_root), |
| "edit_img_path": rel(edit_img, clean_root), |
| "ori_voxel_num": ori_voxel_num, |
| "edit_voxel_num": edit_voxel_num, |
| } |
|
|
| stats["records"] = len(records) |
| stats["encoded_without_condition_key"] = len(encoded_keys - condition_keys) |
| stats["condition_without_encoded_key"] = len(condition_keys - encoded_keys) |
| return records, stats |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| clean_root = args.clean_root.resolve() |
| output = args.output.resolve() if args.output is not None else clean_root / "dataset_info.json" |
|
|
| if not clean_root.is_dir(): |
| raise FileNotFoundError(f"clean root not found: {clean_root}") |
|
|
| flux, flux_stats = build_branch(clean_root, "flux", "flux_encode", args.voxel_count_mode) |
| alpaca, alpaca_stats = build_branch(clean_root, "alpaca", "alpaca_encode", args.voxel_count_mode) |
|
|
| info = { |
| "mixamo": {}, |
| "flux_edit": flux, |
| "alpaca": alpaca, |
| "_meta": { |
| "clean_root": str(clean_root), |
| "source": "generated_from_3Deditverse_data", |
| "voxel_count_mode": args.voxel_count_mode, |
| "branches": { |
| "flux_edit": flux_stats, |
| "alpaca": alpaca_stats, |
| }, |
| "total_editverse_records": len(flux) + len(alpaca), |
| }, |
| } |
|
|
| output.parent.mkdir(parents=True, exist_ok=True) |
| indent = None if args.indent == 0 else args.indent |
| output.write_text(json.dumps(info, indent=indent, sort_keys=True) + "\n") |
|
|
| print(f"Wrote: {output}") |
| print(f"Flux records: {len(flux)}") |
| print(f"Alpaca records: {len(alpaca)}") |
| print(f"Total records: {len(flux) + len(alpaca)}") |
| print("Stats:") |
| print(json.dumps(info["_meta"]["branches"], indent=2, sort_keys=True)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
