| 261 | |
| 262 | |
| 263 | class ReidTestAugmentation(object): |
| 264 | def __init__(self, height, width, vit=False): |
| 265 | |
| 266 | normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], |
| 267 | std=[0.229, 0.224, 0.225]) |
| 268 | |
| 269 | if not vit: |
| 270 | self.test_transformer = T.Compose([ |
| 271 | T.Resize((height, width)), |
| 272 | T.ToTensor(), |
| 273 | normalizer, |
| 274 | ]) |
| 275 | else: |
| 276 | self.test_transformer = T.Compose([ |
| 277 | T.Resize((height, width)), |
| 278 | T.PILToTensor(), |
| 279 | ]) |
| 280 | |
| 281 | def __call__(self, img): |
| 282 | ## transform |
| 283 | return self.test_transformer(img) |
| 284 | |
| 285 | |
| 286 | |