File size: 5,071 Bytes
90ca518 | 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 | """多分辨率网格生成 —— 端到端推理。
低分辨率生成器(TSSR 三前向离散扩散) 逐级生成到格 64, 每级 20 步;
res32 / res64 两级产物各补一次洞再喂下一级;
高分辨率细分器继续 64 -> 128 -> 256 -> 512, 每级 20 步, 全程补洞。
用法:
python scripts/infer.py --data <mesh 目录> --out_dir outputs/run1
python scripts/infer.py --data <mesh 目录> --out_dir outputs/run1 --ncase 20 --empty_bias 0.3
"""
import argparse
import os
import sys
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
DEFAULTS = dict(
STAGE="2,4,8,16,32,64", # 生成器级联阶梯(产物真实体素格)
STEPS="20", # 生成器每级扩散步数
SUPERRES="1", SUPERRES_TO="512", SUPERRES_STEPS="20", # 细分器 64->512, 每级步数
FILL="1", # 级联内补洞(res32/res64, 见 FILL_MIN_RES)
SR_FILL_IN="1", SUPERRES_FIX_HOLES="1", # 细分入口 + 每级输出补洞
DECIMATE_RATIO="0.5", # 输入预塌缩比(0=不塌缩)
METHODS="plain", MULTI_PC="0",
SAVE_GT_LEVELS="1", SAVE_PC_LEVELS="1",
CONFIG=os.path.join(ROOT, "configs", "generator.yaml"),
SUPERRES_CONFIG=os.path.join(ROOT, "configs", "subdivider.yaml"),
)
def main():
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--data", required=True, help="输入 mesh 目录(glb/obj/ply 混合)")
ap.add_argument("--out_dir", default=os.path.join(ROOT, "outputs"), help="输出根目录")
ap.add_argument("--ckpt", default=os.path.join(ROOT, "ckpt", "generator.pth"), help="生成器权重")
ap.add_argument("--subdivider_ckpt", default=os.path.join(ROOT, "ckpt", "subdivider.pth"), help="细分器权重")
ap.add_argument("--out_suffix", default="", help="输出子目录后缀")
ap.add_argument("--ncase", type=int, default=0, help="只跑前 N 个(文件名字母序); 0=全部")
ap.add_argument("--case", default="", help="只跑指定 case, 逗号分隔文件名子串(建议带尾点: 'a.' 不会匹配 'a1.glb')")
ap.add_argument("--gpu", default=None, help="CUDA_VISIBLE_DEVICES")
ap.add_argument("--steps", type=int, default=None, help="覆盖每级步数(默认 20)")
ap.add_argument("--no_fill", action="store_true", help="关闭所有补洞")
ap.add_argument("--fill_min_res", type=int, default=0, help=">0: 级联内补洞只在产物格>=该值的级做(默认 0=res32 起)")
ap.add_argument("--empty_bias", type=float, default=0.0, help=">0: 连接头空类 logit 偏置, 稀疏化控面数(0.3 约 -25%% 面)")
a = ap.parse_args()
for k, v in DEFAULTS.items():
os.environ.setdefault(k, v)
os.environ["RAIL_CKPT"] = os.path.abspath(a.ckpt)
os.environ["SUPERRES_CKPT"] = os.path.abspath(a.subdivider_ckpt)
os.environ["BASE_DIR"] = os.path.dirname(os.path.abspath(a.ckpt))
os.environ["OUT_ROOT"] = os.path.abspath(a.out_dir)
os.environ["DATASET_DIR"] = os.path.abspath(a.data)
os.environ.setdefault("DATASET", "custom")
os.environ["NCASE"] = str(a.ncase if a.ncase > 0 else 10 ** 9)
if a.case: os.environ["CASE_STEM"] = a.case
if a.out_suffix: os.environ["OUT_SUFFIX"] = a.out_suffix
if a.gpu: os.environ["CUDA_VISIBLE_DEVICES"] = a.gpu
if a.steps: os.environ["STEPS"] = os.environ["SUPERRES_STEPS"] = str(a.steps)
if a.no_fill: os.environ.update(FILL="0", SR_FILL_IN="0", SUPERRES_FIX_HOLES="0")
if a.fill_min_res: os.environ["FILL_MIN_RES"] = str(a.fill_min_res)
if a.empty_bias: os.environ["EMPTY_BIAS"] = str(a.empty_bias)
os.chdir(ROOT)
import torch
import modules.pipeline as P
from models.generator import FactorizedTSSRTransformer
class TSSRAdapter(FactorizedTSSRTransformer):
"""把主管线的 .sample() 转发到三前向 sample_tssr; forward 缺省补 mode=1(mask 步)。"""
@torch.no_grad()
def sample(self, coord_list, depth, pc=None, pc_target=None, T_infer=10, cfg_scale=1.0,
temperature=0.0, cfg_schedule="constant", cfg_start=1.0, **_ignored):
return self.sample_tssr(coord_list, depth, pc=pc, pc_target=pc_target, T_infer=T_infer,
cfg_scale=cfg_scale, cfg_schedule=cfg_schedule, cfg_start=cfg_start,
mask_logits_lo=float(os.environ.get("MASK_LO", "0.2")),
mask_logits_hi=float(os.environ.get("MASK_HI", "0.9")))
def forward_varlen(self, *args, **kwargs):
kwargs.pop("cond_cache", None)
if kwargs.get("mode", None) is None and len(args) < 7:
kwargs["mode"] = 1
return super().forward_varlen(*args, **kwargs)
P.BaseModel = TSSRAdapter
P.main()
if __name__ == "__main__":
main()
|