| 106 | |
| 107 | |
| 108 | class Extractor(object): |
| 109 | def __init__(self, model_path, use_cuda=True): |
| 110 | self.net = Net(reid=True) |
| 111 | self.device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu" |
| 112 | state_dict = torch.load(model_path, map_location=torch.device(self.device))[ |
| 113 | 'net_dict'] |
| 114 | self.net.load_state_dict(state_dict) |
| 115 | logger = logging.getLogger("root.tracker") |
| 116 | logger.info("Loading weights from {}... Done!".format(model_path)) |
| 117 | self.net.to(self.device) |
| 118 | self.size = (64, 128) |
| 119 | self.norm = transforms.Compose([ |
| 120 | transforms.ToTensor(), |
| 121 | transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), |
| 122 | ]) |
| 123 | |
| 124 | def _preprocess(self, im_crops): |
| 125 | """ |
| 126 | TODO: |
| 127 | 1. to float with scale from 0 to 1 |
| 128 | 2. resize to (64, 128) as Market1501 dataset did |
| 129 | 3. concatenate to a numpy array |
| 130 | 3. to torch Tensor |
| 131 | 4. normalize |
| 132 | """ |
| 133 | def _resize(im, size): |
| 134 | return cv2.resize(im.astype(np.float32)/255., size) |
| 135 | |
| 136 | im_batch = torch.cat([self.norm(_resize(im, self.size)).unsqueeze( |
| 137 | 0) for im in im_crops], dim=0).float() |
| 138 | return im_batch |
| 139 | |
| 140 | def __call__(self, im_crops): |
| 141 | im_batch = self._preprocess(im_crops) |
| 142 | with torch.no_grad(): |
| 143 | im_batch = im_batch.to(self.device) |
| 144 | features = self.net(im_batch) |
| 145 | return features.cpu().numpy() |