| 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 | """ |