Entry for evaluation.
(self, stage: str = "val", splats_list: Optional[List[Dict]] = None)
| 770 | |
| 771 | @torch.no_grad() |
| 772 | def eval(self, stage: str = "val", splats_list: Optional[List[Dict]] = None): |
| 773 | """Entry for evaluation.""" |
| 774 | print("Running evaluation...") |
| 775 | cfg = self.cfg |
| 776 | device = self.device |
| 777 | world_rank = self.world_rank |
| 778 | world_size = self.world_size |
| 779 | |
| 780 | # dict to save metrics of each frame |
| 781 | seq_stats = defaultdict(dict) |
| 782 | |
| 783 | # if splats_list is not provided, use the default splats_list |
| 784 | if splats_list is None: |
| 785 | splats_list_to_render = self.splats_list |
| 786 | else: |
| 787 | splats_list_to_render = splats_list |
| 788 | if not isinstance(splats_list_to_render[0], torch.nn.ParameterDict): |
| 789 | splats_list_to_render = [ |
| 790 | torch.nn.ParameterDict({ |
| 791 | k: torch.nn.Parameter(v) if isinstance(v, torch.Tensor) else v |
| 792 | for k, v in splats.items() |
| 793 | }) |
| 794 | for splats in splats_list_to_render |
| 795 | ] |
| 796 | |
| 797 | # loop on frame |
| 798 | for f_id, (splats, val_dataset, train_dataset) in enumerate(zip(splats_list_to_render, self.valset_list, self.trainset_list)): |
| 799 | valloader = torch.utils.data.DataLoader( |
| 800 | val_dataset, batch_size=1, shuffle=False, num_workers=1 |
| 801 | ) |
| 802 | ellipse_time = 0 |
| 803 | metrics = defaultdict(list) |
| 804 | # loop on view |
| 805 | for v_id, data in enumerate(valloader): |
| 806 | camtoworlds = data["camtoworld"].to(device) |
| 807 | Ks = data["K"].to(device) |
| 808 | pixels = data["image"].to(device) / 255.0 |
| 809 | masks = data["mask"].to(device) if "mask" in data else None |
| 810 | height, width = pixels.shape[1:3] |
| 811 | splats = splats.to(device) |
| 812 | |
| 813 | torch.cuda.synchronize() |
| 814 | tic = time.time() |
| 815 | colors, _, _ = self.rasterize_splats( |
| 816 | camtoworlds=camtoworlds, |
| 817 | Ks=Ks, |
| 818 | width=width, |
| 819 | height=height, |
| 820 | sh_degree=cfg.sh_degree, |
| 821 | near_plane=cfg.near_plane, |
| 822 | far_plane=cfg.far_plane, |
| 823 | masks=masks, |
| 824 | splats=splats # must need |
| 825 | ) # [1, H, W, 3] |
| 826 | torch.cuda.synchronize() |
| 827 | ellipse_time += time.time() - tic |
| 828 | |
| 829 | colors = torch.clamp(colors, 0.0, 1.0) |