| 353 | return images |
| 354 | |
| 355 | class Stack(object): |
| 356 | |
| 357 | def __init__(self, roll=False): |
| 358 | self.roll = roll |
| 359 | |
| 360 | def __call__(self, img_group): |
| 361 | if img_group[0].mode == 'L': |
| 362 | return np.concatenate([np.expand_dims(x, 2) |
| 363 | for x in img_group], axis=2) |
| 364 | elif img_group[0].mode == 'RGB': |
| 365 | if self.roll: |
| 366 | return np.concatenate([np.array(x)[:, :, ::-1] |
| 367 | for x in img_group], axis=2) |
| 368 | else: |
| 369 | # print(np.concatenate(img_group, axis=2).shape) |
| 370 | # print(img_group[0].shape) |
| 371 | return np.concatenate(img_group, axis=2) |
| 372 | |
| 373 | class ToTorchFormatTensor(object): |
| 374 | """ Converts a PIL.Image (RGB) or numpy.ndarray (H x W x C) in the range [0, 255] |
no outgoing calls