| """多分辨率网格生成 —— 端到端推理。 |
| |
| 低分辨率生成器(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", |
| FILL="1", |
| SR_FILL_IN="1", SUPERRES_FIX_HOLES="1", |
| DECIMATE_RATIO="0.5", |
| 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() |
|
|
