| 8 | |
| 9 | |
| 10 | class BucketFactory: |
| 11 | |
| 12 | def __init__( |
| 13 | self, |
| 14 | ori_thw_list: List[Tuple], |
| 15 | dp_size: int, |
| 16 | rnd_state: np.random.RandomState, |
| 17 | bucket_config: BucketConfig, |
| 18 | ) -> None: |
| 19 | self.ori_thw_list = ori_thw_list |
| 20 | self.dp_size = dp_size |
| 21 | self.rnd_state = rnd_state |
| 22 | self.bucket_config = bucket_config |
| 23 | |
| 24 | def __call__(self, ) -> Any: |
| 25 | flop_list = [ |
| 26 | self.bucket_config(*thw, self.rnd_state) |
| 27 | for thw in self.ori_thw_list |
| 28 | ] |
| 29 | |
| 30 | sorted_indices = sorted( |
| 31 | range(len(flop_list)), key=lambda x: flop_list[x]) |
| 32 | |
| 33 | bucket_dict = OrderedDict({}) |
| 34 | for i, idx in enumerate(sorted_indices): |
| 35 | # video_size = flop_list[idx].size |
| 36 | bucket_key = flop_list[idx] |
| 37 | if bucket_key not in bucket_dict: |
| 38 | bucket_dict[bucket_key] = [] |
| 39 | bucket_dict[bucket_key].append(idx) |
| 40 | |
| 41 | def merge_bucket(bucket_dict: OrderedDict, min_length: int): |
| 42 | all_pass = False |
| 43 | while not all_pass: |
| 44 | all_pass = True |
| 45 | keys = list(bucket_dict.keys()) |
| 46 | to_pop_keys = [] |
| 47 | for key_idx, (key, idx_list) in enumerate(bucket_dict.items()): |
| 48 | if len(idx_list) < min_length: |
| 49 | all_pass = False |
| 50 | if key_idx == len(bucket_dict) - 1: |
| 51 | continue |
| 52 | tgt_key = keys[key_idx + 1] |
| 53 | bucket_dict[tgt_key] = idx_list + bucket_dict[tgt_key] |
| 54 | to_pop_keys.append(key) |
| 55 | for key in to_pop_keys: |
| 56 | bucket_dict.pop(key) |
| 57 | keys = list(bucket_dict.keys()) |
| 58 | if len(bucket_dict[keys[-1]]) < min_length: |
| 59 | bucket_dict[keys[-2]] = bucket_dict[ |
| 60 | keys[-2]] + bucket_dict[keys[-1]] |
| 61 | bucket_dict.pop(keys[-1]) |
| 62 | return bucket_dict |
| 63 | |
| 64 | for bucket_key, idx_list in bucket_dict.items(): |
| 65 | print( |
| 66 | f'before merging, bucket_key: {bucket_key}, consists of {len(idx_list)} clips' |
| 67 | ) |