Method stack

lada/models/basicvsrpp/mmagic/data_sample.py:277–321

Stack a list of data samples to one. All tensor fields will be stacked at first dimension. Otherwise the values will be saved in a list. Args: data_samples (Sequence['DataSample']): A sequence of `DataSample` to stack. Returns:

(cls, data_samples: Sequence['DataSample'])

    @classmethod
    def stack(cls, data_samples: Sequence['DataSample']) -> 'DataSample':
        """Stack a list of data samples to one. All tensor fields will be
        stacked at first dimension. Otherwise the values will be saved in a
        list.

        Args:
            data_samples (Sequence['DataSample']): A sequence of
                `DataSample` to stack.

        Returns:
            DataSample: The stacked data sample.
        """
        # 1. check key consistency
        keys = data_samples[0].keys()
        assert all([data.keys() == keys for data in data_samples])

        meta_keys = data_samples[0].metainfo_keys()
        assert all(
            [data.metainfo_keys() == meta_keys for data in data_samples])

        # 2. stack data
        stacked_data_sample = DataSample()
        for k in keys:
            values = [getattr(data, k) for data in data_samples]
            # 3. check type consistent
            value_type = type(values[0])
            assert all([type(val) == value_type for val in values])

            # 4. stack
            if isinstance(values[0], torch.Tensor):
                stacked_value = torch.stack(values)
            elif isinstance(values[0], LabelData):
                labels = [data.label for data in values]
                values = torch.stack(labels)
                stacked_value = LabelData(label=values)
            else:
                stacked_value = values
            stacked_data_sample.set_field(stacked_value, k)

        # 5. stack metainfo
        for k in meta_keys:
            values = [data.metainfo[k] for data in data_samples]
            stacked_data_sample.set_metainfo({k: values})

        return stacked_data_sample

    def split(self,
              allow_nonseq_value: bool = False) -> Sequence['DataSample']: