LTT/PRM

1import json2import os3import torch4import psutil5import gc6from tqdm import tqdm7from concurrent.futures import ThreadPoolExecutor, as_completed8from src.data.objaverse import load_obj9from src.utils import mesh10from src.utils.material import Material11import argparse12 13 14def bytes_to_megabytes(bytes):15    return bytes / (1024 * 1024)16 17 18def bytes_to_gigabytes(bytes):19    return bytes / (1024 * 1024 * 1024)20 21 22def print_memory_usage(stage):23    process = psutil.Process(os.getpid())24    memory_info = process.memory_info()25    allocated = torch.cuda.memory_allocated() / 1024**226    cached = torch.cuda.memory_reserved() / 1024**227    print(28        f"[{stage}] Process memory: {memory_info.rss / 1024**2:.2f} MB, "29        f"Allocated CUDA memory: {allocated:.2f} MB, Cached CUDA memory: {cached:.2f} MB"30    )31 32 33def process_obj(index, root_dir, final_save_dir, paths):34    obj_path = os.path.join(root_dir, paths[index], paths[index] + '.obj')35    mtl_path = os.path.join(root_dir, paths[index], paths[index] + '.mtl')36 37    if os.path.exists(os.path.join(final_save_dir, f"{paths[index]}.pth")):38        return None39 40    try:41        with torch.no_grad():42            ref_mesh, vertices, faces, normals, nfaces, texcoords, tfaces, uber_material = load_obj(43                obj_path, return_attributes=True44            )45            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")46            ref_mesh = mesh.compute_tangents(ref_mesh)47 48        with open(mtl_path, 'r') as file:49            lines = file.readlines()50 51        if len(lines) >= 250:52            return None53 54        final_mesh_attributes = {55            "v_pos": ref_mesh.v_pos.detach().cpu(),56            "v_nrm": ref_mesh.v_nrm.detach().cpu(),57            "v_tex": ref_mesh.v_tex.detach().cpu(),58            "v_tng": ref_mesh.v_tng.detach().cpu(),59            "t_pos_idx": ref_mesh.t_pos_idx.detach().cpu(),60            "t_nrm_idx": ref_mesh.t_nrm_idx.detach().cpu(),61            "t_tex_idx": ref_mesh.t_tex_idx.detach().cpu(),62            "t_tng_idx": ref_mesh.t_tng_idx.detach().cpu(),63            "mat_dict": {key: ref_mesh.material[key] for key in ref_mesh.material.mat_keys},64        }65 66        torch.save(final_mesh_attributes, f"{final_save_dir}/{paths[index]}.pth")67        print(f"==> Saved to {final_save_dir}/{paths[index]}.pth")68 69        del ref_mesh70        torch.cuda.empty_cache()71        return paths[index]72 73    except Exception as e:74        print(f"Failed to process {paths[index]}: {e}")75        return None76 77    finally:78        gc.collect()79        torch.cuda.empty_cache()80 81 82def main(root_dir, save_dir):83    os.makedirs(save_dir, exist_ok=True)84    finish_lists = os.listdir(save_dir)85    paths = os.listdir(root_dir)86 87    valid_uid = []88 89    print_memory_usage("Start")90 91    batch_size = 10092    num_batches = (len(paths) + batch_size - 1) // batch_size93 94    for batch in tqdm(range(num_batches)):95        start_index = batch * batch_size96        end_index = min(start_index + batch_size, len(paths))97 98        with ThreadPoolExecutor(max_workers=8) as executor:99            futures = [100                executor.submit(process_obj, index, root_dir, save_dir, paths)101                for index in range(start_index, end_index)102            ]103            for future in as_completed(futures):104                result = future.result()105                if result is not None:106                    valid_uid.append(result)107 108        print_memory_usage(f"=====> After processing batch {batch + 1}")109        torch.cuda.empty_cache()110        gc.collect()111 112    print_memory_usage("End")113 114 115if __name__ == "__main__":116    parser = argparse.ArgumentParser(description="Process OBJ files and save final results.")117    parser.add_argument("root_dir", type=str, help="Directory containing the root OBJ files.")118    parser.add_argument("save_dir", type=str, help="Directory to save the processed results.")119    args = parser.parse_args()120 121    main(args.root_dir, args.save_dir)122