Spaces:
Paused
Paused
File size: 6,195 Bytes
08b23ce |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 |
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import numpy as np
import h5py
from tqdm import tqdm
import scipy.sparse as sp
from concurrent.futures import ProcessPoolExecutor, as_completed
from functools import partial
from dataclasses import dataclass
from typing import List, Dict, Any, Optional
from utils.util import read_obj_file, read_rig_file, normalize_to_unit_cube, build_adjacency_list, compute_graph_distance, get_tpl_edges
@dataclass
class ProcessedSample:
"""Data structure for a processed sample."""
vertices: np.ndarray
faces: np.ndarray
joints: np.ndarray
bones: np.ndarray
root_index: int
pc_w_norm: np.ndarray
file_name: str
skin: np.ndarray
graph_dist: np.ndarray
edges: np.ndarray
def process_sample(data: Dict[str, Any]) -> Optional[ProcessedSample]:
"""
Process a single sample from the dataset.
Args:
data: Dictionary containing sample data
Returns:
ProcessedSample object or None if processing fails
"""
vertices = data['vertices'].copy()
joints = data['joints'].copy()
if len(joints) > 70: # filter out data with too many joints
return None
vertices, center, scale = normalize_to_unit_cube(vertices, 0.9995)
joints = (joints - center) * scale
# Build skinning weights matrix
skinning_data = data['skinning_weights_value']
skinning_rows = data['skinning_weights_row']
skinning_cols = data['skinning_weights_col']
skinning_shape = data['skinning_weights_shape']
skinning_sparse = sp.coo_matrix(
(skinning_data, (skinning_rows, skinning_cols)),
shape=skinning_shape
)
skinning_weights = skinning_sparse.toarray() # (n_vertex, n_joints)
# Compute topology and graph features
edges = get_tpl_edges(data['vertices'], data['faces'])
num_joints = len(data['joints'])
adjacency = build_adjacency_list(num_joints, data['bones'])
graph_dist = compute_graph_distance(num_joints, adjacency)
return ProcessedSample(
vertices=vertices,
faces=data['faces'],
joints=joints,
bones=data['bones'],
root_index=data['root_index'],
pc_w_norm=data['pc_w_norm'],
file_name=data['uuid'],
skin=skinning_weights,
graph_dist=graph_dist,
edges=edges
)
def parallel_process_samples(
data_list: List[Dict[str, Any]],
max_workers: Optional[int] = None
) -> List[ProcessedSample]:
"""
Process multiple samples in parallel.
Args:
data_list: List of sample dictionaries
max_workers: Maximum number of worker processes
Returns:
List of successfully processed samples
"""
processed_samples = []
with ProcessPoolExecutor(max_workers=max_workers) as executor:
# Submit all tasks
futures = {executor.submit(process_sample, data): data for data in data_list}
# Process results with progress bar
for future in tqdm(as_completed(futures), total=len(futures), desc='Processing samples'):
try:
result = future.result()
if result is not None:
processed_samples.append(result)
else:
original_data = futures[future]
except Exception as e:
original_data = futures[future]
print(f"Exception in processing {original_data.get('file_name', 'unknown')}: {e}")
return processed_samples
def save_to_h5(processed_samples: List[ProcessedSample], output_path: str) -> None:
"""
Save processed samples to HDF5 file.
Args:
processed_samples: List of processed samples
output_path: Output HDF5 file path
"""
with h5py.File(output_path, 'w') as f:
# Add metadata
f.attrs['num_samples'] = len(processed_samples)
f.attrs['version'] = '1.0'
for i, sample in enumerate(tqdm(processed_samples, desc='Saving to HDF5')):
grp = f.create_group(f'sample_{i}')
# Save arrays with compression
grp.create_dataset('joints', data=sample.joints, compression='gzip')
grp.create_dataset('bones', data=sample.bones, compression='gzip')
grp.create_dataset('root_index', data=sample.root_index, dtype='i')
grp.create_dataset('pc_w_norm', data=sample.pc_w_norm, compression='gzip')
grp.create_dataset('vertices', data=sample.vertices, compression='gzip')
grp.create_dataset('faces', data=sample.faces, compression='gzip')
grp.create_dataset('edges', data=sample.edges, compression='gzip')
grp.create_dataset('skin', data=sample.skin, compression='gzip')
grp.create_dataset('graph_dist', data=sample.graph_dist, compression='gzip')
string_dtype = h5py.special_dtype(vlen=str)
grp.create_dataset('file_name', data=sample.file_name, dtype=string_dtype)
def main(npz_file_path, h5_file_path, max_workers):
loaded_data = np.load(npz_file_path, allow_pickle=True)
data_list = loaded_data['arr_0']
num_samples = len(data_list)
print(f"Total samples: {num_samples}")
processed_samples = parallel_process_samples(
data_list=data_list,
max_workers=max_workers
)
save_to_h5(processed_samples, h5_file_path)
print("Processing complete!")
if __name__ == '__main__':
npz_file_path = 'articulation_xlv2_test.npz'
h5_file_path = 'articulation_xlv2_test.h5'
main(npz_file_path, h5_file_path, max_workers=8)
|