| 221 | return bucket_class(**class_kwargs) |
| 222 | |
| 223 | def __init__(self, ) -> None: |
| 224 | ar_value_arr = np.array(list(self.bucket_config.keys())) |
| 225 | assert np.all(ar_value_arr[:-1] < ar_value_arr[1:]) |
| 226 | self.ar_value_arr = ar_value_arr |
| 227 | for ar, (hwp, prob) in self.bucket_config.items(): |
| 228 | assert isinstance(hwp, np.ndarray) |
| 229 | assert isinstance(prob, np.ndarray) |
| 230 | if hwp.shape[1] == 2: |
| 231 | hwp = np.stack([hwp[:, 0], hwp[:, 1], hwp[:, 0] * hwp[:, 1]], |
| 232 | axis=-1) |
| 233 | else: |
| 234 | assert np.all(hwp[:-1, 2] <= hwp[1:, 2]) |
| 235 | if len(hwp) != len(prob): |
| 236 | raise ValueError( |
| 237 | f'wrong config of aspect ratio: {ar}, size_list: {hwp}, prob_list: {prob}' |
| 238 | ) |
| 239 | prob[:] = prob[:] / prob.sum() |
| 240 | self.bucket_config[ar] = (hwp, prob) |
| 241 | print(f'aspect ratio: {ar}, size list: {self.bucket_config[ar]}') |
| 242 | pass |
| 243 | |
| 244 | def preprocess(self, n_frame: int, height: int, width: int, |
| 245 | rnd_state: np.random.RandomState): |