| 3 | from gloss import * |
| 4 | |
| 5 | class ViewerCallback(Callback): |
| 6 | |
| 7 | def __init__(self, viewer_config_path, experiment_name): |
| 8 | gloss_setup_logger(log_level=LogLevel.Info) |
| 9 | self.viewer=Viewer(viewer_config_path) |
| 10 | self.scene=self.viewer.get_scene() |
| 11 | |
| 12 | self.first_time=True |
| 13 | |
| 14 | self.visualize_every_x_iters=100 |
| 15 | self.visualize_strand_at_idx=None |
| 16 | # self.visualize_strand_at_idx=0 |
| 17 | |
| 18 | def after_forward_pass(self, phase, gt_cloud=None, pred_cloud=None, **kwargs): |
| 19 | |
| 20 | # if phase.iter_nr==1000: |
| 21 | # exit(1) |
| 22 | |
| 23 | if phase.iter_nr%self.visualize_every_x_iters==0: |
| 24 | if self.visualize_strand_at_idx is not None: |
| 25 | gt_cloud=gt_cloud[self.visualize_strand_at_idx,:,:] |
| 26 | gt_ent = self.scene.get_or_spawn_renderable("gt_ent") |
| 27 | gt_ent.insert(Verts(gt_cloud.reshape(-1,3).cpu().numpy())) |
| 28 | gt_ent.insert(Colors(gt_cloud.reshape(-1,3).cpu().numpy())) #just to avoid verts being different than Colors when we switch nr_verts for the same entity |
| 29 | if self.first_time: |
| 30 | gt_ent.insert(VisPoints(show_points=True, \ |
| 31 | point_size=3.0, \ |
| 32 | color_type=PointColorType.Solid)) |
| 33 | |
| 34 | |
| 35 | #pred |
| 36 | if self.visualize_strand_at_idx is not None: |
| 37 | pred_cloud=pred_cloud[self.visualize_strand_at_idx,:,:] |
| 38 | pred_ent = self.scene.get_or_spawn_renderable("pred_ent") |
| 39 | pred_ent.insert(Verts(pred_cloud.reshape(-1,3).cpu().numpy())) |
| 40 | pred_ent.insert(Colors(pred_cloud.reshape(-1,3).cpu().numpy())) #just to avoid verts being different than Colors when we switch nr_verts for the same entity |
| 41 | if self.first_time: |
| 42 | pred_ent.insert(VisPoints(show_points=True, \ |
| 43 | point_size=3.0, \ |
| 44 | point_color=[0.1, 0.2, 0.8, 1.0], |
| 45 | color_type=PointColorType.Solid)) |
| 46 | |
| 47 | #render |
| 48 | self.viewer.start_frame() |
| 49 | self.viewer.update() |
| 50 | |
| 51 | |
| 52 | self.first_time=False |
| 53 | |
| 54 | |
| 55 | |