update inference

inference.py

@@ -85,6 +85,21 @@ def resize_foreground(
 def extract_texmesh(args, model, output_path, device):
     # Prepare directory
     ins_dir = output_path
+    # Noise Filter
+    srt_param = model.srt_param.clone()
+    prim_position = srt_param[:, 1:4]
+    prim_scale = srt_param[:, 0:1]
+    dist = torch.sqrt(torch.sum((prim_position[:, None, :] - prim_position[None, :, :]) ** 2, dim=-1))
+    dist += torch.eye(prim_position.shape[0]).to(srt_param)
+    min_dist, min_indices = dist.min(1)
+    dst_prim_scale = prim_scale[min_indices, :]
+    min_scale_converage = prim_scale * 1.414 + dst_prim_scale * 1.414
+    prim_mask = min_dist < min_scale_converage[:, 0]
+    filtered_srt_param = srt_param[prim_mask, :]
+    filtered_feat_param = model.feat_param.clone()[prim_mask, ...]
+    model.srt_param.data = filtered_srt_param
+    model.feat_param.data = filtered_feat_param
+    print(f'[INFO] Mesh Extraction on PrimX: srt={model.srt_param.shape} feat={model.feat_param.shape}')
 
     # Get SDFs
     with torch.no_grad():
@@ -350,9 +365,6 @@ if __name__ == "__main__":
     # manually enable tf32 to get speedup on A100 GPUs
     torch.backends.cuda.matmul.allow_tf32 = True
     torch.backends.cudnn.allow_tf32 = True
-    os.environ["CC"] = "/mnt/lustre/share/gcc/gcc-8.5.0/bin/gcc"
-    os.environ["CPP"] = "/mnt/lustre/share/gcc/gcc-8.5.0/bin/g++"
-    os.environ["CXX"] = "/mnt/lustre/share/gcc/gcc-8.5.0/bin/g++"
     # set config
     config = OmegaConf.load(str(sys.argv[1]))
     config_cli = OmegaConf.from_cli(args_list=sys.argv[2:])