| 222 | |
| 223 | |
| 224 | class ReidTestAugmentationCV2(object): |
| 225 | def __init__(self, height, width, vit=False): |
| 226 | self.height = height |
| 227 | self.width = width |
| 228 | self.normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], |
| 229 | std=[0.229, 0.224, 0.225]) |
| 230 | self.vit = vit |
| 231 | |
| 232 | def __call__(self, img): |
| 233 | ## transform |
| 234 | img = img[:,:,::-1] |
| 235 | img = cv2.resize(img, (self.width, self.height)) |
| 236 | img = img.transpose((2,0,1)) |
| 237 | if not self.vit: |
| 238 | img = torch.Tensor(img/255.) |
| 239 | img = self.normalizer(img) |
| 240 | else: |
| 241 | img = torch.Tensor(img) |
| 242 | return img |
| 243 | |
| 244 | |
| 245 | ## Author: liuyakun1 |