Spaces:
Running
on
L40S
Running
on
L40S
File size: 2,505 Bytes
2252f3d |
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 |
import os
import json
import math
from concurrent.futures import ProcessPoolExecutor
import threading
from tqdm import tqdm
# from glcontext import egl
# egl.create_context()
# exit(0)
LOCAL_RANK = 0
num_processes = 4
NODE_RANK = int(os.getenv("SLURM_PROCID"))
WORLD_SIZE = 1
NODE_NUM=1
# NODE_RANK = int(os.getenv("SLURM_NODEID"))
IS_MAIN = False
if NODE_RANK == 0 and LOCAL_RANK == 0:
IS_MAIN = True
GLOBAL_RANK = NODE_RANK * (WORLD_SIZE//NODE_NUM) + LOCAL_RANK
# json_path = "object_lists/Thuman2.0.json"
# json_path = "object_lists/THuman3.0.json"
json_path = "object_lists/CustomHumans.json"
data_dir = '/aifs4su/mmcode/lipeng'
save_dir = '/aifs4su/mmcode/lipeng/human_8view_new'
def parse_obj_list(x):
if 'THuman3.0' in x:
# print(apath)
splits = x.split('/')
x = os.path.join('THuman3.0', splits[-2])
elif 'Thuman2.0' in x:
splits = x.split('/')
x = os.path.join('Thuman2.0', splits[-2])
elif 'CustomHumans' in x:
splits = x.split('/')
x = os.path.join('CustomHumans', splits[-2])
# print(splits[-2])
elif '1M' in x:
splits = x.split('/')
x = os.path.join('2K2K', splits[-2])
elif 'realistic_8k_model' in x:
splits = x.split('/')
x = os.path.join('realistic_8k_model', splits[-1].split('.')[0])
return f'{save_dir}/{x}'
with open(json_path, 'r') as f:
glb_list = json.load(f)
# glb_list = ['Thuman2.0/0011/0011.obj']
# glb_list = ['THuman3.0/00024_1/00024_0006/mesh.obj']
# glb_list = ['CustomHumans/mesh/0383_00070_02_00061/mesh-f00061.obj']
# glb_list = ['realistic_8k_model/1d41f2a72f994306b80e632f1cc8233f.glb']
total_num_glbs = len(glb_list)
num_glbs_local = int(math.ceil(total_num_glbs / WORLD_SIZE))
start_idx = GLOBAL_RANK * num_glbs_local
end_idx = start_idx + num_glbs_local
# print(start_idx, end_idx)
local_glbs = glb_list[start_idx:end_idx]
if IS_MAIN:
pbar = tqdm(total=len(local_glbs))
lock = threading.Lock()
def process_human(glb_path):
src_path = os.path.join(data_dir, glb_path)
save_path = parse_obj_list(glb_path)
# print(save_path)
command = ('blender -b -P blender_render_human_script.py'
f' -- --object_path {src_path}'
f' --output_dir {save_path} ')
# 1>/dev/null
# print(command)
os.system(command)
if IS_MAIN:
with lock:
pbar.update(1)
with ProcessPoolExecutor(max_workers=num_processes) as executor:
executor.map(process_human, local_glbs)
|