| 254 | |
| 255 | |
| 256 | class Stack(object): |
| 257 | |
| 258 | def __init__(self, roll=False): |
| 259 | self.roll = roll |
| 260 | |
| 261 | def __call__(self, img_group): |
| 262 | mode = img_group[0].mode |
| 263 | if mode == '1': |
| 264 | img_group = [img.convert('L') for img in img_group] |
| 265 | mode = 'L' |
| 266 | if mode == 'L': |
| 267 | return np.stack([np.expand_dims(x, 2) for x in img_group], axis=2) |
| 268 | elif mode == 'RGB': |
| 269 | if self.roll: |
| 270 | return np.stack([np.array(x)[:, :, ::-1] for x in img_group], axis=2) |
| 271 | else: |
| 272 | return np.stack(img_group, axis=2) |
| 273 | else: |
| 274 | raise NotImplementedError(f"Image mode {mode}") |
| 275 | |
| 276 | |
| 277 | class ToTorchFormatTensor(object): |
no outgoing calls
no test coverage detected