Spaces:
Running
on
Zero
Running
on
Zero
File size: 6,357 Bytes
0f079b2 |
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 172 173 174 175 |
# multiprocessing render
import json
import multiprocessing
import subprocess
from dataclasses import dataclass
from typing import Optional
import os
import boto3
import argparse
parser = argparse.ArgumentParser(description='distributed rendering')
parser.add_argument('--workers_per_gpu', type=int,
help='number of workers per gpu.')
parser.add_argument('--input_models_path', type=str,
help='Path to a json file containing a list of 3D object files.')
parser.add_argument('--upload_to_s3', type=bool, default=False,
help='Whether to upload the rendered images to S3.')
parser.add_argument('--log_to_wandb', type=bool, default=False,
help='Whether to log the progress to wandb.')
parser.add_argument('--num_gpus', type=int, default=-1,
help='number of gpus to use. -1 means all available gpus.')
parser.add_argument('--gpu_list',nargs='+', type=int,
help='the avalaible gpus')
parser.add_argument('--mode', type=str, default='render',
choices=['render_ortho', 'render_persp'],
help='use orthogonal camera or perspective camera')
parser.add_argument('--start_i', type=int, default=0,
help='the index of first object to be rendered.')
parser.add_argument('--end_i', type=int, default=-1,
help='the index of the last object to be rendered.')
parser.add_argument('--objaverse_root', type=str, default='/ghome/l5/xxlong/.objaverse/hf-objaverse-v1',
help='Path to a json file containing a list of 3D object files.')
parser.add_argument('--save_folder', type=str, default=None,
help='Path to a json file containing a list of 3D object files.')
parser.add_argument('--blender_install_path', type=str, default=None,
help='blender path.')
parser.add_argument('--view_idx', type=int, default=2,
help='the number of render views.')
parser.add_argument('--ortho_scale', type=float, default=1.25,
help='ortho rendering usage; how large the object is')
parser.add_argument('--random_pose', action='store_true',
help='whether randomly rotate the poses to be rendered')
args = parser.parse_args()
view_idx = args.view_idx
VIEWS = ["front", "back", "right", "left", "front_right", "front_left", "back_right", "back_left"]
def check_task_finish(render_dir, view_index):
files_type = ['rgb', 'normals']
flag = True
view_index = "%03d" % view_index
if os.path.exists(render_dir):
for t in files_type:
for face in VIEWS:
fpath = os.path.join(render_dir, f'{t}_{view_index}_{face}.webp')
# print(fpath)
if not os.path.exists(fpath):
flag = False
else:
flag = False
return flag
def worker(
queue: multiprocessing.JoinableQueue,
count: multiprocessing.Value,
gpu: int,
s3: Optional[boto3.client],
) -> None:
while True:
item = queue.get()
if item is None:
break
view_path = os.path.join(args.save_folder, item.split('/')[-1][:2], item.split('/')[-1][:-4])
print(view_path)
if 'render' in args.mode:
if check_task_finish(view_path, view_idx):
queue.task_done()
print('========', item, 'rendered', '========')
continue
else:
os.makedirs(view_path, exist_ok = True)
# Perform some operation on the item
print(item, gpu)
if args.mode == 'render_ortho':
command = (
f" CUDA_VISIBLE_DEVICES={gpu} "
f" blenderproc run --blender-install-path {args.blender_install_path} blenderProc_ortho.py"
f" --object_path {item} --view {view_idx}"
f" --output_folder {args.save_folder}"
f" --ortho_scale {args.ortho_scale} "
)
if args.random_pose:
print("random pose to render")
command += f" --random_pose"
elif args.mode == 'render_persp':
command = (
f" CUDA_VISIBLE_DEVICES={gpu} "
f" blenderproc run --blender-install-path {args.blender_install_path} blenderProc_persp.py"
f" --object_path {item} --view {view_idx}"
f" --output_folder {args.save_folder}"
)
if args.random_pose:
print("random pose to render")
command += f" --random_pose"
print(command)
subprocess.run(command, shell=True)
with count.get_lock():
count.value += 1
queue.task_done()
if __name__ == "__main__":
# args = tyro.cli(Args)
s3 = boto3.client("s3") if args.upload_to_s3 else None
queue = multiprocessing.JoinableQueue()
count = multiprocessing.Value("i", 0)
# Start worker processes on each of the GPUs
for gpu_i in range(args.num_gpus):
for worker_i in range(args.workers_per_gpu):
worker_i = gpu_i * args.workers_per_gpu + worker_i
process = multiprocessing.Process(
target=worker, args=(queue, count, args.gpu_list[gpu_i], s3)
)
process.daemon = True
process.start()
# Add items to the queue
if args.input_models_path is not None:
with open(args.input_models_path, "r") as f:
model_paths = json.load(f)
args.end_i = len(model_paths) if args.end_i > len(model_paths) else args.end_i
for item in model_paths[args.start_i:args.end_i]:
if os.path.exists(os.path.join(args.objaverse_root, os.path.basename(item))):
obj_path = os.path.join(args.objaverse_root, os.path.basename(item))
elif os.path.exists(os.path.join(args.objaverse_root, item)):
obj_path = os.path.join(args.objaverse_root, item)
else:
obj_path = os.path.join(args.objaverse_root, item[:2], item+".glb")
queue.put(obj_path)
# Wait for all tasks to be completed
queue.join()
# Add sentinels to the queue to stop the worker processes
for i in range(args.num_gpus * args.workers_per_gpu):
queue.put(None)
|