| 89 | |
| 90 | |
| 91 | def assign_probability(bucket_config, prob_dict): |
| 92 | for ar, size_list in bucket_config.items(): |
| 93 | size_list = sorted(size_list, key=lambda x: x[2]) |
| 94 | size_list_with_prob = [] |
| 95 | for pixel_range, prob in prob_dict.items(): |
| 96 | in_range_size_list = [ |
| 97 | size for size in size_list |
| 98 | if size[0] * size[1] >= pixel_range[0] and size[0] * |
| 99 | size[1] < pixel_range[1] |
| 100 | ] |
| 101 | if len(in_range_size_list) == 0: |
| 102 | continue |
| 103 | in_range_size_list = [(*size, prob) for size in in_range_size_list] |
| 104 | size_list_with_prob = size_list_with_prob + in_range_size_list |
| 105 | bucket_config[ar] = (np.array(size_list_with_prob)[:, :2].astype( |
| 106 | np.int64), np.array(size_list_with_prob)[:, -1]) |
| 107 | bucket_config[ar][ |
| 108 | 1][:] = bucket_config[ar][1][:] / bucket_config[ar][1][:].sum() |
| 109 | return bucket_config |
| 110 | |
| 111 | |
| 112 | class FlopEstimator(ABC): |