Update src/pixel3dmm/tracking/tracker.py

src/pixel3dmm/tracking/tracker.py

@@ -1279,6 +1279,18 @@ class Tracker(object):
                                      verts_depth=proj_vertices[:, :, 2:3],
                                      is_viz=True
                                      )
+            # if they asked *only* for the pure shape mask:
+            if visualizations == [[View.SHAPE]]:
+                # ops['normal_images'] is [1,3,H,W] in world‐space ∈ [-1,1]
+                normals = ops['normal_images'][0].cpu().numpy()     # [3,H,W]
+                # remap to [0,1]
+                normals = (normals + 1.0) / 2.0
+                # H×W×3
+                normals = np.transpose(normals, (1, 2, 0))
+                # scale to uint8
+                arr = (normals * 255).clip(0,255).astype(np.uint8)
+                return arr
+
             mask = (self.parse_mask(ops, batch, visualization=True) > 0).float()
             grabbed_depth = ops['actual_rendered_depth'][0, 0,
             torch.clamp(proj_vertices[0, :, 1].long(), 0, self.config.size-1),
@@ -1289,49 +1301,6 @@ class Tracker(object):
                 is_visible_verts_idx = torch.ones_like(is_visible_verts_idx)
 
 
-            all_final_views = []
-            for b_i in range(bs):
-                final_views = []
-
-                for views in visualizations:
-                    row = []
-                    for view in views:
-                        if view == View.COLOR_OVERLAY:
-                            row.append((ops['normal_images'][b_i].cpu().numpy() + 1)/2)
-                        if view == View.GROUND_TRUTH:
-                            row.append(images[b_i].cpu().numpy())
-                        if (view == View.LANDMARKS and not self.no_lm) or is_camera:
-                            gt_lmks = images[b_i:b_i+1].clone()
-                            gt_lmks = util.tensor_vis_landmarks(gt_lmks, landmarks[b_i:b_i+1, :, :], color='g')
-                            gt_lmks = util.tensor_vis_landmarks(gt_lmks, batch['left_iris'][b_i:b_i+1, ...], color='g')
-                            gt_lmks = util.tensor_vis_landmarks(gt_lmks, batch['right_iris'][b_i:b_i+1, ...], color='g')
-                            gt_lmks = util.tensor_vis_landmarks(gt_lmks, proj_vertices[b_i:b_i+1, left_iris_flame, ...], color='r')
-                            gt_lmks = util.tensor_vis_landmarks(gt_lmks, proj_vertices[b_i:b_i+1, right_iris_flame, ...], color='r')
-                            gt_lmks = util.tensor_vis_landmarks(gt_lmks, lmk68[b_i:b_i+1, :, :], color='r')
-                            row.append(gt_lmks[0].cpu().numpy())
-
-                    if True:
-                        nvd_mask = gaussian_blur(ops['mask_images_rendering'].detach(),
-                                                 kernel_size=[self.config.normal_mask_ksize, self.config.normal_mask_ksize],
-                                                 sigma=[self.config.normal_mask_ksize, self.config.normal_mask_ksize])
-                        nvd_mask = (nvd_mask > 0.5).float()
-                        nvd_mask_clone = nvd_mask.clone()
-
-
-                        eyebrow_level = torch.min(lmk68[:, :, 1], dim=1).indices
-
-                        for _i in range(eyebrow_level.shape[0]):
-                            nvd_mask_clone[_i, :, :eyebrow_level[_i], :] = 0
-
-
-                    final_views.append(row)
-
-
-                # VIDEO
-                final_views = util.merge_views(final_views)
-                all_final_views.append(final_views)
-            final_views = np.concatenate(all_final_views, axis=0)
-
             if outer_iter is None:
                 frame_id = str(self.frame).zfill(5)
             else:
@@ -1714,8 +1683,9 @@ class Tracker(object):
             batches = {k: torch.cat([x[k] for x in batches], dim=0) for k in batch.keys()}
             selected_frames = torch.from_numpy(np.array(selected_frames)).long().cuda()
 
-            result_rendering = self.render_and_save(batches, visualizations=[[View.SHAPE]],
-                                                    frame_dst='/joint_initialization', outer_iter=0, timestep=timestep, is_final=True, selected_frames=selected_frames)
+            result_rendering = self.render_and_save(batch,
+                                                    visualizations=[[View.SHAPE]],   # ← only mesh by default
+                                                    frame_dst='/video', save=True, dump_directly=False, outer_iter=0, timestep=timestep, is_final=True, selected_frames=selected_frames)
             video_frames.append(np.array(result_rendering))
             self.frame += 1