add back normal map visualization

apps/infer.py

@@ -269,26 +269,37 @@ def generate_model(in_path, model_type):
         os.makedirs(os.path.join(config_dict['out_dir'],
                     cfg.name, "obj"), exist_ok=True)
 
-        norm_pred = (
+        norm_pred_F = (
             ((in_tensor["normal_F"][0].permute(1, 2, 0) + 1.0) * 255.0 / 2.0)
             .detach()
             .cpu()
             .numpy()
             .astype(np.uint8)
         )
+        
+        norm_pred_B = (
+            ((in_tensor["normal_B"][0].permute(1, 2, 0) + 1.0) * 255.0 / 2.0)
+            .detach()
+            .cpu()
+            .numpy()
+            .astype(np.uint8)
+        )
 
-        norm_orig = unwrap(norm_pred, data)
+        norm_orig_F = unwrap(norm_pred_F, data)
+        norm_orig_B = unwrap(norm_pred_B, data)
+        
         mask_orig = unwrap(
             np.repeat(
                 data["mask"].permute(1, 2, 0).detach().cpu().numpy(), 3, axis=2
             ).astype(np.uint8),
             data,
         )
-        rgb_norm = blend_rgb_norm(data["ori_image"], norm_orig, mask_orig)
+        rgb_norm_F = blend_rgb_norm(data["ori_image"], norm_orig_F, mask_orig)
+        rgb_norm_B = blend_rgb_norm(data["ori_image"], norm_orig_B, mask_orig)
 
         Image.fromarray(
             np.concatenate(
-                [data["ori_image"].astype(np.uint8), rgb_norm], axis=1)
+                [data["ori_image"].astype(np.uint8), rgb_norm_F, rgb_norm_B], axis=1)
         ).save(os.path.join(config_dict['out_dir'], cfg.name, f"png/{data['name']}_overlap.png"))
 
         smpl_obj = trimesh.Trimesh(
@@ -448,7 +459,7 @@ def generate_model(in_path, model_type):
         dataset.render.load_meshes(
             verts_lst, faces_lst)
         dataset.render.get_rendered_video(
-            [data["ori_image"], rgb_norm],
+            [data["ori_image"], rgb_norm_F, rgb_norm_B],
             os.path.join(config_dict['out_dir'], cfg.name,
                          f"vid/{data['name']}_cloth.mp4"),
         )