| |
| """Refresh metadata.csv status columns for Edit3D-Bench TRELLIS preprocessing.""" |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import pandas as pd |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--output_dir", required=True) |
| parser.add_argument("--feat_model", default="dinov2_vitl14_reg") |
| parser.add_argument("--ss_latent_name", default="ss_enc_conv3d_16l8_fp16") |
| parser.add_argument("--latent_name", default="dinov2_vitl14_reg_slat_enc_swin8_B_64l8_fp16") |
| args = parser.parse_args() |
|
|
| output_dir = Path(args.output_dir) |
| metadata_path = output_dir / "metadata.csv" |
| df = pd.read_csv(metadata_path) |
|
|
| rendered = [] |
| voxelized = [] |
| num_voxels = [] |
| feature_col = f"feature_{args.feat_model}" |
| ss_col = f"ss_latent_{args.ss_latent_name}" |
| latent_col = f"latent_{args.latent_name}" |
| feature_ok = [] |
| ss_ok = [] |
| latent_ok = [] |
|
|
| for row in df.to_dict("records"): |
| sha = row["sha256"] |
| render_dir = output_dir / "renders" / sha |
| voxel_path = output_dir / "voxels" / f"{sha}.ply" |
| feature_path = output_dir / "features" / args.feat_model / f"{sha}.npz" |
| ss_path = output_dir / "ss_latents" / args.ss_latent_name / f"{sha}.npz" |
| latent_path = output_dir / "latents" / args.latent_name / f"{sha}.npz" |
| rendered.append((render_dir / "transforms.json").exists() and (render_dir / "mesh.ply").exists()) |
| voxelized.append(voxel_path.exists()) |
| feature_ok.append(feature_path.exists()) |
| ss_ok.append(ss_path.exists()) |
| latent_ok.append(latent_path.exists()) |
| if voxel_path.exists(): |
| try: |
| import utils3d |
|
|
| num_voxels.append(int(len(utils3d.io.read_ply(str(voxel_path))[0]))) |
| except Exception: |
| num_voxels.append(0) |
| else: |
| num_voxels.append(0) |
|
|
| df["rendered"] = rendered |
| df["voxelized"] = voxelized |
| df["num_voxels"] = num_voxels |
| df[feature_col] = feature_ok |
| df[ss_col] = ss_ok |
| df[latent_col] = latent_ok |
| df.to_csv(metadata_path, index=False) |
|
|
| print(f"metadata: {metadata_path}") |
| print(f"rendered {sum(rendered)}/{len(df)}") |
| print(f"voxelized {sum(voxelized)}/{len(df)}") |
| print(f"{feature_col} {sum(feature_ok)}/{len(df)}") |
| print(f"{ss_col} {sum(ss_ok)}/{len(df)}") |
| print(f"{latent_col} {sum(latent_ok)}/{len(df)}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
