File size: 2,932 Bytes
434b0b0 | 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 | # -*- coding: utf-8 -*-
# @Organization : Tongyi Lab, Alibaba
# @Author : Lingteng Qiu
# @Email : 220019047@link.cuhk.edu.cn
# @Time : 2025-08-31 10:02:15
# @Function : Distributed MVS rendering pipeline
import sys
sys.path.append(".")
import os
import time
import traceback
import numpy as np
from tqdm import tqdm
from engine.MVSRender.mvs_render import MVSRender
def basename(path):
pre_name = os.path.basename(path).split(".")[0]
return pre_name
def multi_process(worker, items, **kwargs):
"""
worker worker function to process items
"""
nodes = kwargs["nodes"]
dirs = kwargs["dirs"]
bucket = int(
np.ceil(len(items) / nodes)
) # avoid last node process too many items.
print("Total Nodes:", nodes)
print("Save Path:", dirs)
rank = int(os.environ.get("RANK", 0))
print("Current Rank:", rank)
kwargs["RANK"] = rank
if rank == nodes - 1:
output_dir = worker(items[bucket * rank :], **kwargs)
else:
output_dir = worker(items[bucket * rank : bucket * (rank + 1)], **kwargs)
if rank == 0 and nodes > 1:
sleep_time = int(kwargs.get("timesleep", 3600))
time.sleep(sleep_time) # one hour
def run_mvsrender(items, **params):
output_dir = params["dirs"]
debug = params["debug"]
os.makedirs(output_dir, exist_ok=True)
if debug:
items = items[:5]
process_valid = []
for item in tqdm(items, desc="Processing..."):
if ".ply" not in item:
continue
ply_basename = basename(item)
rendering_save_path = os.path.join(output_dir, ply_basename)
if not debug and os.path.exists(rendering_save_path):
continue
try:
# load ply
renderer = MVSRender(item, radius=2.5) # Radius is the rendering distance.
# rendering ply
renderer.rendering(rendering_save_path)
print(f"rendering: {rendering_save_path}")
except:
traceback.print_exc()
process_valid.append(item)
return output_dir
def get_parse():
import argparse
parser = argparse.ArgumentParser(description="")
parser.add_argument("-i", "--input", required=True, help="input path")
parser.add_argument("-o", "--output", required=True, help="output path")
parser.add_argument("--nodes", default=1, type=int, help="how many workload?")
parser.add_argument("--debug", action="store_true", help="debug tag")
args = parser.parse_args()
return args
if __name__ == "__main__":
opt = get_parse()
# catch avaliable items
available_items = os.listdir(opt.input)
available_items = [os.path.join(opt.input, item) for item in available_items]
multi_process(
worker=run_mvsrender,
items=available_items,
dirs=opt.output,
nodes=opt.nodes,
debug=opt.debug,
)
|