Method __call__

datasets/multi_resolution_sampler.py:24–89

(self, )

        self.bucket_config = bucket_config

    def __call__(self, ) -> Any:
        flop_list = [
            self.bucket_config(*thw, self.rnd_state)
            for thw in self.ori_thw_list
        ]

        sorted_indices = sorted(
            range(len(flop_list)), key=lambda x: flop_list[x])

        bucket_dict = OrderedDict({})
        for i, idx in enumerate(sorted_indices):
            # video_size = flop_list[idx].size
            bucket_key = flop_list[idx]
            if bucket_key not in bucket_dict:
                bucket_dict[bucket_key] = []
            bucket_dict[bucket_key].append(idx)

        def merge_bucket(bucket_dict: OrderedDict, min_length: int):
            all_pass = False
            while not all_pass:
                all_pass = True
                keys = list(bucket_dict.keys())
                to_pop_keys = []
                for key_idx, (key, idx_list) in enumerate(bucket_dict.items()):
                    if len(idx_list) < min_length:
                        all_pass = False
                        if key_idx == len(bucket_dict) - 1:
                            continue
                        tgt_key = keys[key_idx + 1]
                        bucket_dict[tgt_key] = idx_list + bucket_dict[tgt_key]
                        to_pop_keys.append(key)
                for key in to_pop_keys:
                    bucket_dict.pop(key)
                keys = list(bucket_dict.keys())
                if len(bucket_dict[keys[-1]]) < min_length:
                    bucket_dict[keys[-2]] = bucket_dict[
                        keys[-2]] + bucket_dict[keys[-1]]
                    bucket_dict.pop(keys[-1])
            return bucket_dict

        for bucket_key, idx_list in bucket_dict.items():
            print(
                f'before merging, bucket_key: {bucket_key}, consists of {len(idx_list)} clips'
            )

        bucket_dict = merge_bucket(bucket_dict, min_length=self.dp_size)
        for bucket_key, idx_list in bucket_dict.items():
            print(
                f'after merging, bucket_key: {bucket_key}, consists of {len(idx_list)} clips'
            )
        final_idx_list = []
        final_bucket_key_list = []

        for video_size, idx_list in bucket_dict.items():
            print(
                f'before appending, bucket_key: {video_size}, consists of {len(idx_list)} clips'
            )
            if (res := (len(idx_list) % self.dp_size)) != 0: