| import json |
| import os |
| import torch |
| import psutil |
| import gc |
| from tqdm import tqdm |
| from concurrent.futures import ThreadPoolExecutor, as_completed |
| from src.data.objaverse import load_obj |
| from src.utils import mesh |
| from src.utils.material import Material |
| import argparse |
|
|
|
|
| def bytes_to_megabytes(bytes): |
| return bytes / (1024 * 1024) |
|
|
|
|
| def bytes_to_gigabytes(bytes): |
| return bytes / (1024 * 1024 * 1024) |
|
|
|
|
| def print_memory_usage(stage): |
| process = psutil.Process(os.getpid()) |
| memory_info = process.memory_info() |
| allocated = torch.cuda.memory_allocated() / 1024**2 |
| cached = torch.cuda.memory_reserved() / 1024**2 |
| print( |
| f"[{stage}] Process memory: {memory_info.rss / 1024**2:.2f} MB, " |
| f"Allocated CUDA memory: {allocated:.2f} MB, Cached CUDA memory: {cached:.2f} MB" |
| ) |
|
|
|
|
| def process_obj(index, root_dir, final_save_dir, paths): |
| obj_path = os.path.join(root_dir, paths[index], paths[index] + '.obj') |
| mtl_path = os.path.join(root_dir, paths[index], paths[index] + '.mtl') |
|
|
| if os.path.exists(os.path.join(final_save_dir, f"{paths[index]}.pth")): |
| return None |
|
|
| try: |
| with torch.no_grad(): |
| ref_mesh, vertices, faces, normals, nfaces, texcoords, tfaces, uber_material = load_obj( |
| obj_path, return_attributes=True |
| ) |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| ref_mesh = mesh.compute_tangents(ref_mesh) |
|
|
| with open(mtl_path, 'r') as file: |
| lines = file.readlines() |
|
|
| if len(lines) >= 250: |
| return None |
|
|
| final_mesh_attributes = { |
| "v_pos": ref_mesh.v_pos.detach().cpu(), |
| "v_nrm": ref_mesh.v_nrm.detach().cpu(), |
| "v_tex": ref_mesh.v_tex.detach().cpu(), |
| "v_tng": ref_mesh.v_tng.detach().cpu(), |
| "t_pos_idx": ref_mesh.t_pos_idx.detach().cpu(), |
| "t_nrm_idx": ref_mesh.t_nrm_idx.detach().cpu(), |
| "t_tex_idx": ref_mesh.t_tex_idx.detach().cpu(), |
| "t_tng_idx": ref_mesh.t_tng_idx.detach().cpu(), |
| "mat_dict": {key: ref_mesh.material[key] for key in ref_mesh.material.mat_keys}, |
| } |
|
|
| torch.save(final_mesh_attributes, f"{final_save_dir}/{paths[index]}.pth") |
| print(f"==> Saved to {final_save_dir}/{paths[index]}.pth") |
|
|
| del ref_mesh |
| torch.cuda.empty_cache() |
| return paths[index] |
|
|
| except Exception as e: |
| print(f"Failed to process {paths[index]}: {e}") |
| return None |
|
|
| finally: |
| gc.collect() |
| torch.cuda.empty_cache() |
|
|
|
|
| def main(root_dir, save_dir): |
| os.makedirs(save_dir, exist_ok=True) |
| finish_lists = os.listdir(save_dir) |
| paths = os.listdir(root_dir) |
|
|
| valid_uid = [] |
|
|
| print_memory_usage("Start") |
|
|
| batch_size = 100 |
| num_batches = (len(paths) + batch_size - 1) // batch_size |
|
|
| for batch in tqdm(range(num_batches)): |
| start_index = batch * batch_size |
| end_index = min(start_index + batch_size, len(paths)) |
|
|
| with ThreadPoolExecutor(max_workers=8) as executor: |
| futures = [ |
| executor.submit(process_obj, index, root_dir, save_dir, paths) |
| for index in range(start_index, end_index) |
| ] |
| for future in as_completed(futures): |
| result = future.result() |
| if result is not None: |
| valid_uid.append(result) |
|
|
| print_memory_usage(f"=====> After processing batch {batch + 1}") |
| torch.cuda.empty_cache() |
| gc.collect() |
|
|
| print_memory_usage("End") |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description="Process OBJ files and save final results.") |
| parser.add_argument("root_dir", type=str, help="Directory containing the root OBJ files.") |
| parser.add_argument("save_dir", type=str, help="Directory to save the processed results.") |
| args = parser.parse_args() |
|
|
| main(args.root_dir, args.save_dir) |
|
|
