diff --git a/087_DeepSort/download.sh b/087_DeepSort/download.sh new file mode 100755 index 0000000000..1648c4c412 --- /dev/null +++ b/087_DeepSort/download.sh @@ -0,0 +1,9 @@ +#!/bin/bash + +curl -sc /tmp/cookie "https://drive.google.com/uc?export=download&id=1CIAU0OQL4bFIms-e74V2OUYnqIzq10q-" > /dev/null +CODE="$(awk '/_warning_/ {print $NF}' /tmp/cookie)" +curl -Lb /tmp/cookie "https://drive.google.com/uc?export=download&confirm=${CODE}&id=1CIAU0OQL4bFIms-e74V2OUYnqIzq10q-" -o resources.tar.gz +tar -zxvf resources.tar.gz +rm resources.tar.gz + +echo Download finished. diff --git a/087_DeepSort/feature_extractor.py b/087_DeepSort/feature_extractor.py new file mode 100644 index 0000000000..7e46c5830b --- /dev/null +++ b/087_DeepSort/feature_extractor.py @@ -0,0 +1,60 @@ +import torch +import torchvision.transforms as transforms +import numpy as np +import cv2 +import logging + +from .model import Net + + +class Extractor(object): + def __init__(self, model_path, use_cuda=True): + self.net = Net(reid=True) + self.device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu" + print('@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@ model_path:', model_path) + state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)['net_dict'] + self.net.load_state_dict(state_dict) + + logger = logging.getLogger("root.tracker") + logger.info("Loading weights from {}... Done!".format(model_path)) + self.net.to(self.device) + self.size = (64, 128) + self.norm = transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), + ]) + + def _preprocess(self, im_crops): + """ + TODO: + 1. to float with scale from 0 to 1 + 2. resize to (64, 128) as Market1501 dataset did + 3. concatenate to a numpy array + 3. to torch Tensor + 4. normalize + """ + def _resize(im, size): + return cv2.resize(im.astype(np.float32)/255., size) + + im_batch = torch.cat([self.norm(_resize(im, self.size)).unsqueeze( + 0) for im in im_crops], dim=0).float() + return im_batch + + def __call__(self, im_crops): + im_batch = self._preprocess(im_crops) + with torch.no_grad(): + im_batch = im_batch.to(self.device) + features = self.net(im_batch) + + ################################################################################################ + torch.onnx.export(self.net, im_batch, 'deepsort_128x64.onnx', verbose=True, opset_version=12) + ################################################################################################ + + return features.cpu().numpy() + + +if __name__ == '__main__': + img = cv2.imread("demo.jpg")[:, :, (2, 1, 0)] + extr = Extractor("checkpoint/ckpt.t7") + feature = extr(img) + print(feature.shape) diff --git a/087_DeepSort/url.txt b/087_DeepSort/url.txt new file mode 100644 index 0000000000..ff193dde01 --- /dev/null +++ b/087_DeepSort/url.txt @@ -0,0 +1,2 @@ +https://github.com/mikel-brostrom/Yolov5_DeepSort_Pytorch +https://github.com/PINTO0309/openvino2tensorflow \ No newline at end of file