wehub-resource-sync
597 lines
24 KiB
Python
597 lines
24 KiB
Python
"""
|
|
voxelize_pbr_view.py - Multi-view transform PBR voxelization
|
|
Extends voxelize_pbr.py with scale and mesh rotation logic
|
|
Based on dual_grid_view.py and test_ovoxel_pbr_transform.py implementation
|
|
"""
|
|
import os
|
|
import copy
|
|
import sys
|
|
import importlib
|
|
import argparse
|
|
import json
|
|
import math
|
|
import pandas as pd
|
|
import pickle
|
|
import numpy as np
|
|
import torch
|
|
from easydict import EasyDict as edict
|
|
from functools import partial
|
|
import o_voxel
|
|
from utils import get_new_camera_matrix, sphere_normalize_torch
|
|
|
|
|
|
# ==================== PBR-specific transform functions ====================
|
|
|
|
def transform_vertices(vertices, frame):
|
|
"""
|
|
Apply multi-view transform to vertices based on camera transform matrix.
|
|
|
|
Args:
|
|
vertices: torch.Tensor, shape [N, 3], vertex coordinates
|
|
frame: dict containing transform_matrix
|
|
|
|
Returns:
|
|
transformed_vertices: torch.Tensor, shape [N, 3]
|
|
"""
|
|
device = vertices.device
|
|
c2w_orig = torch.tensor(frame['transform_matrix'], dtype=torch.float32, device=device)
|
|
|
|
# Old and new camera matrices
|
|
radius = c2w_orig[:3, 3].norm().item()
|
|
c2w_new = get_new_camera_matrix(radius=radius, yaw=-90/180.0*math.pi, pitch=0.0,
|
|
dtype=torch.float32, device=device)
|
|
w2c_orig = torch.inverse(c2w_orig)
|
|
|
|
# Initial and final axis alignment matrices
|
|
R_init = torch.tensor([
|
|
[1.0, 0.0, 0.0, 0.0],
|
|
[0.0, 0.0, -1.0, 0.0],
|
|
[0.0, 1.0, 0.0, 0.0],
|
|
[0.0, 0.0, 0.0, 1.0]
|
|
], dtype=torch.float32, device=device)
|
|
|
|
R_back = torch.tensor([
|
|
[1.0, 0.0, 0.0, 0.0],
|
|
[0.0, 0.0, 1.0, 0.0],
|
|
[0.0, -1.0, 0.0, 0.0],
|
|
[0.0, 0.0, 0.0, 1.0]
|
|
], dtype=torch.float32, device=device)
|
|
|
|
R_ply = torch.tensor([
|
|
[1.0, 0.0, 0.0, 0.0],
|
|
[0.0, 0.0, 1.0, 0.0],
|
|
[0.0, -1.0, 0.0, 0.0],
|
|
[0.0, 0.0, 0.0, 1.0]
|
|
], dtype=torch.float32, device=device)
|
|
|
|
T_cam = c2w_new @ w2c_orig @ R_ply
|
|
T_final = R_back @ T_cam @ R_init
|
|
|
|
# Apply transform
|
|
vertices = vertices.reshape(-1, 3)
|
|
verts_h = torch.cat([vertices, torch.ones((vertices.shape[0], 1), dtype=torch.float32, device=device)], dim=1)
|
|
verts_trans = (T_final @ verts_h.T).T[:, :3]
|
|
|
|
return verts_trans
|
|
|
|
|
|
def transform_normals(normals, frame):
|
|
"""
|
|
Apply multi-view transform to normals (rotation only).
|
|
Consistent with test_ovoxel_pbr_transform.py implementation.
|
|
|
|
Args:
|
|
normals: torch.Tensor or np.ndarray, shape [N, 3] or [N, 3, 3]
|
|
frame: dict containing transform_matrix
|
|
|
|
Returns:
|
|
transformed_normals: np.ndarray (always returns numpy for dump compatibility)
|
|
"""
|
|
is_numpy = isinstance(normals, np.ndarray)
|
|
if is_numpy:
|
|
normals = torch.from_numpy(normals).float()
|
|
|
|
device = normals.device
|
|
original_shape = normals.shape
|
|
|
|
# Flatten to [N, 3] for processing
|
|
if len(original_shape) == 3:
|
|
normals_flat = normals.reshape(-1, 3)
|
|
else:
|
|
normals_flat = normals
|
|
|
|
c2w_orig = torch.tensor(frame['transform_matrix'], dtype=torch.float32, device=device)
|
|
|
|
# Old and new camera matrices
|
|
radius = c2w_orig[:3, 3].norm().item()
|
|
c2w_new = get_new_camera_matrix(radius=radius, yaw=-90/180.0*math.pi, pitch=0.0,
|
|
dtype=torch.float32, device=device)
|
|
w2c_orig = torch.inverse(c2w_orig)
|
|
|
|
# Axis alignment matrices (rotation part only, 3x3)
|
|
R_init = torch.tensor([
|
|
[1.0, 0.0, 0.0],
|
|
[0.0, 0.0, -1.0],
|
|
[0.0, 1.0, 0.0]
|
|
], dtype=torch.float32, device=device)
|
|
|
|
R_back = torch.tensor([
|
|
[1.0, 0.0, 0.0],
|
|
[0.0, 0.0, 1.0],
|
|
[0.0, -1.0, 0.0]
|
|
], dtype=torch.float32, device=device)
|
|
|
|
R_ply = torch.tensor([
|
|
[1.0, 0.0, 0.0],
|
|
[0.0, 0.0, 1.0],
|
|
[0.0, -1.0, 0.0]
|
|
], dtype=torch.float32, device=device)
|
|
|
|
# Use rotation part only
|
|
T_cam_rot = c2w_new[:3, :3] @ w2c_orig[:3, :3] @ R_ply
|
|
T_final_rot = R_back @ T_cam_rot @ R_init
|
|
|
|
# Apply rotation transform
|
|
normals_trans = torch.matmul(normals_flat, T_final_rot.T)
|
|
|
|
# Re-normalize
|
|
normals_trans = torch.nn.functional.normalize(normals_trans, dim=-1)
|
|
|
|
# Restore original shape
|
|
if len(original_shape) == 3:
|
|
normals_trans = normals_trans.reshape(original_shape)
|
|
|
|
# Always return numpy array for dump compatibility
|
|
return normals_trans.numpy()
|
|
|
|
|
|
def prepare_pbr_dump(dump):
|
|
"""
|
|
Prepare PBR dump data for processing.
|
|
Consistent with voxelize_pbr.py preprocessing.
|
|
|
|
Args:
|
|
dump: raw PBR dump data
|
|
|
|
Returns:
|
|
processed dump data
|
|
"""
|
|
dump = copy.deepcopy(dump)
|
|
|
|
# Fix dump alpha map
|
|
for mat in dump['materials']:
|
|
if mat['alphaTexture'] is not None and mat['alphaMode'] == 'OPAQUE':
|
|
mat['alphaMode'] = 'BLEND'
|
|
|
|
# Append default material
|
|
dump['materials'].append({
|
|
"baseColorFactor": [0.8, 0.8, 0.8],
|
|
"alphaFactor": 1.0,
|
|
"metallicFactor": 0.0,
|
|
"roughnessFactor": 0.5,
|
|
"alphaMode": "OPAQUE",
|
|
"alphaCutoff": 0.5,
|
|
"baseColorTexture": None,
|
|
"alphaTexture": None,
|
|
"metallicTexture": None,
|
|
"roughnessTexture": None,
|
|
})
|
|
|
|
# Filter out empty objects
|
|
dump['objects'] = [
|
|
obj for obj in dump['objects']
|
|
if obj['vertices'].size != 0 and obj['faces'].size != 0
|
|
]
|
|
|
|
return dump
|
|
|
|
|
|
def transform_pbr_dump(dump, frame):
|
|
"""
|
|
Apply multi-view transform to entire PBR dump data.
|
|
|
|
Processing flow (based on test_ovoxel_pbr_transform.py):
|
|
1. Box normalize all vertices (scale only, no center shift)
|
|
2. Sphere normalize
|
|
3. Apply multi-view transform
|
|
4. Normalize back to [-0.5, 0.5]^3
|
|
|
|
Note: All object vertices are processed together (not per-object) for consistency.
|
|
|
|
Args:
|
|
dump: PBR dump data (already preprocessed via prepare_pbr_dump)
|
|
frame: camera frame info
|
|
|
|
Returns:
|
|
transformed_dump: transformed dump data
|
|
total_scale: total scale from original mesh to final mesh
|
|
"""
|
|
transformed_dump = copy.deepcopy(dump)
|
|
|
|
# 1. Collect all vertices
|
|
all_vertices_list = []
|
|
vertex_counts = []
|
|
for obj in transformed_dump['objects']:
|
|
all_vertices_list.append(obj['vertices'])
|
|
vertex_counts.append(len(obj['vertices']))
|
|
|
|
if len(all_vertices_list) == 0:
|
|
return transformed_dump, 1.0
|
|
|
|
all_vertices = np.concatenate(all_vertices_list, axis=0)
|
|
|
|
# 2. Box normalize (scale only, no center shift, consistent with original rendering)
|
|
vertices_min = all_vertices.min(axis=0)
|
|
vertices_max = all_vertices.max(axis=0)
|
|
box_scale_init = 0.99999 / (vertices_max - vertices_min).max()
|
|
all_vertices_box_normalized = all_vertices * box_scale_init
|
|
|
|
all_vertices_tensor = torch.from_numpy(all_vertices_box_normalized).float()
|
|
|
|
# 3. Sphere normalize all vertices together
|
|
all_vertices_sphere, sphere_center, sphere_radius = sphere_normalize_torch(all_vertices_tensor)
|
|
|
|
# 4. Multi-view transform
|
|
all_transformed = transform_vertices(all_vertices_sphere, frame)
|
|
|
|
# 5. Normalize back to [-0.5, 0.5]^3 (all vertices together)
|
|
abs_max = all_transformed.abs().max().item()
|
|
box_scale_final = 0.49999 / abs_max
|
|
all_transformed_normalized = all_transformed * box_scale_final
|
|
|
|
# Compute total scale (from original mesh to final normalized mesh)
|
|
total_scale = box_scale_init * box_scale_final / sphere_radius.item()
|
|
|
|
# 6. Split back to individual objects
|
|
start_idx = 0
|
|
for i, obj in enumerate(transformed_dump['objects']):
|
|
end_idx = start_idx + vertex_counts[i]
|
|
obj['vertices'] = all_transformed_normalized[start_idx:end_idx].numpy()
|
|
start_idx = end_idx
|
|
|
|
# Transform normals
|
|
if obj['normals'] is not None and obj['normals'].size > 0:
|
|
obj['normals'] = transform_normals(obj['normals'], frame)
|
|
|
|
# Fix mat_ids (replace -1 with default material index)
|
|
obj['mat_ids'][obj['mat_ids'] == -1] = len(transformed_dump['materials']) - 1
|
|
|
|
# Validate range
|
|
assert np.all(obj['mat_ids'] >= 0), 'invalid mat_ids'
|
|
assert np.all(obj['vertices'] >= -0.5) and np.all(obj['vertices'] <= 0.5), 'vertices out of range'
|
|
|
|
return transformed_dump, total_scale
|
|
|
|
|
|
def _pbr_voxelize_view(file, sha256, pbr_dump_root, transform_root, root, view_indices=None):
|
|
"""
|
|
Process multi-view PBR voxelization for a single sha256.
|
|
|
|
Args:
|
|
file: local_path from metadata
|
|
sha256: sha256 string
|
|
pbr_dump_root: directory containing PBR dump files
|
|
transform_root: directory containing transform json files
|
|
root: output directory for PBR voxels
|
|
view_indices: list of view indices to process, None for all views
|
|
"""
|
|
try:
|
|
pack = {'sha256': sha256}
|
|
dump = None
|
|
|
|
# Load transforms
|
|
transform_path = os.path.join(transform_root, sha256, 'transforms.json')
|
|
if not os.path.exists(transform_path):
|
|
print(f'Transform file not found for {sha256}, skipping')
|
|
return {'sha256': sha256, 'error': 'Transform file not found'}
|
|
|
|
with open(transform_path, 'r') as f:
|
|
transforms_json = json.load(f)
|
|
transform_mats = transforms_json['frames']
|
|
|
|
# Determine views to process
|
|
if view_indices is None:
|
|
view_indices = list(range(len(transform_mats)))
|
|
else:
|
|
view_indices = [i for i in view_indices if i < len(transform_mats)]
|
|
|
|
# Track processed and skipped counts
|
|
processed_count = 0
|
|
skipped_count = 0
|
|
|
|
for view_idx in view_indices:
|
|
for res in opt.resolution:
|
|
need_process = False
|
|
|
|
# Check if already processed
|
|
# Path structure: pbr_voxels_view_fix_{res}/{sha256}/view{idx:02d}.vxz
|
|
sha256_dir = os.path.join(root, f'pbr_voxels_view_fix_{res}', sha256)
|
|
vxz_path = os.path.join(sha256_dir, f'view{view_idx:02d}.vxz')
|
|
if os.path.exists(vxz_path):
|
|
try:
|
|
info = o_voxel.io.read_vxz_info(vxz_path)
|
|
pack[f'pbr_voxelized_view_fix{view_idx:02d}_{res}'] = True
|
|
pack[f'num_pbr_voxels_view_fix{view_idx:02d}_{res}'] = info['num_voxel']
|
|
skipped_count += 1
|
|
except Exception as e:
|
|
print(f'Error reading {sha256}/view{view_idx:02d}.vxz: {e}, will reprocess')
|
|
need_process = True
|
|
else:
|
|
need_process = True
|
|
|
|
# Process PBR dump
|
|
if need_process:
|
|
# Lazy load dump
|
|
if dump is None:
|
|
pbr_dump_file = os.path.join(pbr_dump_root, 'pbr_dumps', f'{sha256}.pickle')
|
|
if not os.path.exists(pbr_dump_file):
|
|
print(f'PBR dump not found for {sha256}, skipping')
|
|
return {'sha256': sha256, 'error': 'PBR dump not found'}
|
|
|
|
with open(pbr_dump_file, 'rb') as f:
|
|
dump = pickle.load(f)
|
|
|
|
# Prepare dump data
|
|
dump = prepare_pbr_dump(dump)
|
|
|
|
if len(dump['objects']) == 0:
|
|
print(f'No valid objects in PBR dump for {sha256}, skipping')
|
|
return {'sha256': sha256, 'error': 'No valid objects in PBR dump'}
|
|
|
|
# Get transform for current view
|
|
frame = transform_mats[view_idx]
|
|
|
|
# Multi-view transform (deep copy from original dump each time)
|
|
transformed_dump, total_scale = transform_pbr_dump(dump, frame)
|
|
|
|
# PBR voxelization
|
|
coord, attr = o_voxel.convert.blender_dump_to_volumetric_attr(
|
|
transformed_dump,
|
|
grid_size=res,
|
|
aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
|
|
mip_level_offset=0,
|
|
verbose=False,
|
|
timing=False
|
|
)
|
|
|
|
# Remove normal and emissive (consistent with voxelize_pbr.py)
|
|
del attr['normal']
|
|
del attr['emissive']
|
|
|
|
# Save .vxz file
|
|
os.makedirs(sha256_dir, exist_ok=True)
|
|
o_voxel.io.write_vxz(vxz_path, coord, attr)
|
|
|
|
# Save scale info
|
|
scale_path = os.path.join(sha256_dir, f'view{view_idx:02d}_scale.json')
|
|
scale_info = {
|
|
'sha256': sha256,
|
|
'view_idx': view_idx,
|
|
'total_scale': float(total_scale),
|
|
}
|
|
with open(scale_path, 'w') as f:
|
|
json.dump(scale_info, f, indent=2)
|
|
|
|
pack[f'pbr_voxelized_view_fix{view_idx:02d}_{res}'] = True
|
|
pack[f'num_pbr_voxels_view_fix{view_idx:02d}_{res}'] = len(coord)
|
|
pack[f'pbr_voxel_scale_view_fix{view_idx:02d}_{res}'] = float(total_scale)
|
|
processed_count += 1
|
|
|
|
# Record processing stats
|
|
pack['_processed_count'] = processed_count
|
|
pack['_skipped_count'] = skipped_count
|
|
|
|
return pack
|
|
|
|
except Exception as e:
|
|
print(f'Error processing {sha256}: {e}')
|
|
import traceback
|
|
traceback.print_exc()
|
|
return {'sha256': sha256, 'error': str(e)}
|
|
|
|
|
|
if __name__ == '__main__':
|
|
dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--root', type=str, required=True,
|
|
help='Directory to save the metadata')
|
|
parser.add_argument('--pbr_dump_root', type=str, default=None,
|
|
help='Directory to load PBR dumps')
|
|
parser.add_argument('--transform_root', type=str, default=None,
|
|
help='Directory to load transform json files (renders_cond)')
|
|
parser.add_argument('--pbr_voxel_root', type=str, default=None,
|
|
help='Directory to save voxelized PBR attributes')
|
|
parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
|
|
help='Filter objects with aesthetic score lower than this value')
|
|
parser.add_argument('--instances', type=str, default=None,
|
|
help='Instances to process')
|
|
parser.add_argument('--view_indices', type=str, default=None,
|
|
help='View indices to process, e.g., "0,1,2" or "0-5". None for all views')
|
|
parser.add_argument('--skip_list', type=str, default=None,
|
|
help='Path to a file containing sha256 hashes to skip (one per line). '
|
|
'Supports format: "sha256" or "dataset/sha256"')
|
|
parser.add_argument('--clean_pbr_dir', type=str, default=None,
|
|
help='Path to clean_pbr directory. Will auto-load {dataset}_clean_output.txt as ok-list, '
|
|
'only sha256 in ok-list will be processed')
|
|
parser.add_argument('--clean_pbr_name', type=str, default=None,
|
|
help='Dataset name prefix for clean_pbr file (e.g., ObjaverseXL_github). '
|
|
'Defaults to sys.argv[1] if not specified')
|
|
dataset_utils.add_args(parser)
|
|
parser.add_argument('--resolution', type=str, default='1024')
|
|
parser.add_argument('--rank', type=int, default=0)
|
|
parser.add_argument('--world_size', type=int, default=1)
|
|
parser.add_argument('--max_workers', type=int, default=0)
|
|
opt = parser.parse_args(sys.argv[2:])
|
|
opt = edict(vars(opt))
|
|
opt.resolution = sorted([int(x) for x in opt.resolution.split(',')], reverse=True)
|
|
opt.pbr_dump_root = opt.pbr_dump_root or opt.root
|
|
opt.transform_root = opt.transform_root or os.path.join(opt.root, 'renders_cond')
|
|
opt.pbr_voxel_root = opt.pbr_voxel_root or opt.root
|
|
|
|
# Parse view_indices
|
|
view_indices = None
|
|
if opt.view_indices is not None:
|
|
view_indices = []
|
|
for part in opt.view_indices.split(','):
|
|
if '-' in part:
|
|
start, end = map(int, part.split('-'))
|
|
view_indices.extend(range(start, end + 1))
|
|
else:
|
|
view_indices.append(int(part))
|
|
view_indices = list(set(view_indices)) # Deduplicate
|
|
view_indices.sort()
|
|
|
|
# Load skip list (sha256 hashes to skip)
|
|
skip_set = set()
|
|
if opt.skip_list is not None and os.path.exists(opt.skip_list):
|
|
with open(opt.skip_list, 'r') as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line and not line.startswith('#'):
|
|
# Support "dataset/sha256" and plain "sha256" format, extract pure sha256
|
|
skip_set.add(line.split('/')[-1])
|
|
print(f'Loaded {len(skip_set)} items from skip_list: {opt.skip_list}')
|
|
|
|
# Load clean_pbr ok-list (only process approved sha256)
|
|
ok_set = None
|
|
if opt.clean_pbr_dir is not None:
|
|
dataset_name = opt.clean_pbr_name or sys.argv[1]
|
|
clean_file = os.path.join(opt.clean_pbr_dir, f'{dataset_name}_clean_output.txt')
|
|
if os.path.exists(clean_file):
|
|
ok_set = set()
|
|
with open(clean_file, 'r') as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line and not line.startswith('#'):
|
|
ok_set.add(line.split('/')[-1])
|
|
print(f'Loaded {len(ok_set)} ok items from clean_pbr: {clean_file}')
|
|
else:
|
|
print(f'Warning: clean_pbr file not found: {clean_file}, proceeding without ok-list filter')
|
|
|
|
for res in opt.resolution:
|
|
os.makedirs(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_view_fix_{res}', 'new_records'), exist_ok=True)
|
|
|
|
# Get file list
|
|
if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
|
|
raise ValueError('metadata.csv not found')
|
|
metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
|
|
if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
|
|
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
|
|
if os.path.exists(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')):
|
|
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')).set_index('sha256'))
|
|
|
|
# Check already processed pbr_voxels_view_fix
|
|
for res in opt.resolution:
|
|
if os.path.exists(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_view_fix_{res}', 'metadata.csv')):
|
|
pbr_voxel_metadata = pd.read_csv(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_view_fix_{res}', 'metadata.csv')).set_index('sha256')
|
|
metadata = metadata.combine_first(pbr_voxel_metadata)
|
|
|
|
metadata = metadata.reset_index()
|
|
|
|
if opt.instances is None:
|
|
if opt.filter_low_aesthetic_score is not None:
|
|
metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
|
|
metadata = metadata[metadata['pbr_dumped'] == True]
|
|
|
|
# Filter out objects with all views already processed
|
|
if view_indices is not None:
|
|
for res in opt.resolution:
|
|
# Check if each specified view is already processed
|
|
all_views_done_col = f'_all_views_done_{res}'
|
|
metadata[all_views_done_col] = True
|
|
for view_idx in view_indices:
|
|
col_name = f'pbr_voxelized_view_fix{view_idx:02d}_{res}'
|
|
if col_name in metadata.columns:
|
|
metadata[all_views_done_col] = metadata[all_views_done_col] & (metadata[col_name] == True)
|
|
else:
|
|
metadata[all_views_done_col] = False
|
|
break
|
|
|
|
# Keep objects with at least one incomplete resolution
|
|
any_incomplete = None
|
|
for res in opt.resolution:
|
|
all_views_done_col = f'_all_views_done_{res}'
|
|
if all_views_done_col in metadata.columns:
|
|
if any_incomplete is None:
|
|
any_incomplete = ~metadata[all_views_done_col]
|
|
else:
|
|
any_incomplete = any_incomplete | ~metadata[all_views_done_col]
|
|
|
|
if any_incomplete is not None:
|
|
before_filter = len(metadata)
|
|
metadata = metadata[any_incomplete]
|
|
print(f'Filtered out {before_filter - len(metadata)} already completed objects')
|
|
else:
|
|
if os.path.exists(opt.instances):
|
|
with open(opt.instances, 'r') as f:
|
|
instances = f.read().splitlines()
|
|
else:
|
|
instances = opt.instances.split(',')
|
|
metadata = metadata[metadata['sha256'].isin(instances)]
|
|
|
|
# Apply skip_list filter (exclude specified sha256)
|
|
if skip_set:
|
|
before_skip = len(metadata)
|
|
metadata = metadata[~metadata['sha256'].isin(skip_set)]
|
|
print(f'Skip list: filtered out {before_skip - len(metadata)} objects, {len(metadata)} remaining')
|
|
|
|
# Apply clean_pbr ok-list filter (only keep approved sha256)
|
|
if ok_set is not None:
|
|
before_ok = len(metadata)
|
|
metadata = metadata[metadata['sha256'].isin(ok_set)]
|
|
print(f'Ok list: kept {len(metadata)} objects out of {before_ok} (filtered {before_ok - len(metadata)})')
|
|
|
|
metadata = metadata.sample(frac=1, random_state=444).reset_index(drop=True)
|
|
start = len(metadata) * opt.rank // opt.world_size
|
|
end = len(metadata) * (opt.rank + 1) // opt.world_size
|
|
metadata = metadata[start:end]
|
|
|
|
print(f'Processing {len(metadata)} objects...')
|
|
if view_indices:
|
|
print(f'View indices to process: {view_indices}')
|
|
else:
|
|
print('Processing all available views')
|
|
|
|
# Process objects
|
|
func = partial(_pbr_voxelize_view,
|
|
pbr_dump_root=opt.pbr_dump_root,
|
|
transform_root=opt.transform_root,
|
|
root=opt.pbr_voxel_root,
|
|
view_indices=view_indices)
|
|
pbr_voxelized = dataset_utils.foreach_instance(metadata, opt.root, func, max_workers=opt.max_workers, desc='Voxelizing PBR views')
|
|
|
|
# Processing summary
|
|
total_processed = pbr_voxelized['_processed_count'].sum() if '_processed_count' in pbr_voxelized.columns else 0
|
|
total_skipped = pbr_voxelized['_skipped_count'].sum() if '_skipped_count' in pbr_voxelized.columns else 0
|
|
print(f'\n========== Processing Summary ==========')
|
|
print(f'Total processed (new): {int(total_processed)}')
|
|
print(f'Total skipped (existing): {int(total_skipped)}')
|
|
print(f'Total items: {int(total_processed + total_skipped)}')
|
|
print(f'=========================================\n')
|
|
|
|
if 'error' in pbr_voxelized.columns:
|
|
errors = pbr_voxelized[pbr_voxelized['error'].notna()]
|
|
if len(errors) > 0:
|
|
with open('errors_pbr_view.txt', 'w') as f:
|
|
f.write('\n'.join(errors['sha256'].tolist()))
|
|
print(f'Errors written to errors_pbr_view.txt ({len(errors)} errors)')
|
|
|
|
# Save metadata
|
|
for res in opt.resolution:
|
|
# Collect all view-related columns
|
|
view_cols = [col for col in pbr_voxelized.columns if f'pbr_voxelized_view_fix' in col and f'_{res}' in col]
|
|
if view_cols:
|
|
# Save metadata for each view
|
|
pbr_voxel_metadata = pbr_voxelized[pbr_voxelized[view_cols].any(axis=1)]
|
|
if len(pbr_voxel_metadata) > 0:
|
|
# Save simplified metadata
|
|
cols_to_save = ['sha256'] + [col for col in pbr_voxelized.columns if f'_{res}' in col]
|
|
cols_to_save = [col for col in cols_to_save if col in pbr_voxelized.columns]
|
|
pbr_voxel_metadata[cols_to_save].to_csv(
|
|
os.path.join(opt.pbr_voxel_root, f'pbr_voxels_view_fix_{res}', 'new_records', f'part_{opt.rank}.csv'),
|
|
index=False
|
|
)
|
|
|
|
print('Done!')
|