Class BucketConfig

datasets/bucket_config.py:211–264

class BucketConfig:
    bucket_config = OrderedDict({})
    flop = FlopEstimatorExact

    @classmethod
    def from_class_name(cls, class_name, class_kwargs: dict):
        bucket_class = globals().get(class_name)
        if not (isinstance(bucket_class, BucketConfig)
                or issubclass(bucket_class, BucketConfig)):
            raise ValueError(f'wrong class name {class_name}')
        return bucket_class(**class_kwargs)

    def __init__(self, ) -> None:
        ar_value_arr = np.array(list(self.bucket_config.keys()))
        assert np.all(ar_value_arr[:-1] < ar_value_arr[1:])
        self.ar_value_arr = ar_value_arr
        for ar, (hwp, prob) in self.bucket_config.items():
            assert isinstance(hwp, np.ndarray)
            assert isinstance(prob, np.ndarray)
            if hwp.shape[1] == 2:
                hwp = np.stack([hwp[:, 0], hwp[:, 1], hwp[:, 0] * hwp[:, 1]],
                               axis=-1)
            else:
                assert np.all(hwp[:-1, 2] <= hwp[1:, 2])
            if len(hwp) != len(prob):
                raise ValueError(
                    f'wrong config of aspect ratio: {ar}, size_list: {hwp}, prob_list: {prob}'
                )
            prob[:] = prob[:] / prob.sum()
            self.bucket_config[ar] = (hwp, prob)
            print(f'aspect ratio: {ar}, size list: {self.bucket_config[ar]}')
        pass

    def preprocess(self, n_frame: int, height: int, width: int,
                   rnd_state: np.random.RandomState):
        """process original height width, sample a (height, width) in
        bucket_config with the closed aspect ratio."""
        ar = height / width
        tgt_ar_idx = find_nearest_value(ar, self.ar_value_arr, return_idx=True)
        h_w_pixel_list, prob_list = self.bucket_config[
            self.ar_value_arr[tgt_ar_idx]]
        if h_w_pixel_list[-1, -1] > height * width:
            down_num = np.sum(h_w_pixel_list[:, -1] <= height * width)
            h_w_pixel_list = h_w_pixel_list[:down_num]
            prob_list = prob_list[:down_num] / prob_list[:down_num].sum()
        selected_idx = rnd_state.choice(
            len(prob_list), size=None, replace=False, p=prob_list)
        h, w = h_w_pixel_list[selected_idx][:2]
        return h, w

    def __call__(self, n_frame: int, height: int, width: int,
                 rnd_state: np.random.RandomState) -> Any:
        tgt_h, tgt_w = self.preprocess(n_frame, height, width, rnd_state)
        return self.flop(n_frame, tgt_h, tgt_w)


class BucketConfigHardCoded1(BucketConfig):