教程 2: 如何设计数据处理流程

@PIPELINES.register_module()
class GetViewdirs:
    """get viewdirs from rays_d
    """
    def __init__(self, enable=True, **kwargs):
        self.enable = enable

    def __call__(self, results):
        """get viewdirs
        Args:
            results (dict): The resulting dict to be modified and passed
                to the next transform in pipeline.
        """
        if self.enable:
            viewdirs = results['rays_d'].clone()
            viewdirs = viewdirs / torch.norm(viewdirs, dim=-1, keepdim=True)
            viewdirs = torch.reshape(viewdirs, [-1, 3]).float()
            results['viewdirs'] = viewdirs
        return results

train_pipeline = [
    dict(type='Sample'),
    dict(type='DeleteUseless', keys=['images', 'poses', 'i_data', 'idx']),
    dict(type='ToTensor', keys=['pose', 'target_s']),
    dict(type='GetRays'),
    dict(type='SelectRays',
        sel_n=N_rand_per_sampler,
        precrop_iters=500,
        precrop_frac=0.5),  # in the first 500 iter, select rays inside center of image
    dict(type='GetViewdirs', enable=use_viewdirs),
    dict(type='ToNDC', enable=(not no_ndc)),
    dict(type='GetBounds'),
    dict(type='GetZvals', lindisp=lindisp,
        N_samples=N_samples),  # N_samples: number of coarse samples per ray
    dict(type='PerturbZvals', enable=is_perturb),
    dict(type='GetPts'),
    dict(type='DeleteUseless', keys=['pose', 'iter_n']),
]

data = {'poses': self.poses, 'images': self.images, 'i_data': self.i_train, 'idx': idx}