由于号称Yolov5_DeepSort_Pytorch之github官网(mikel-brostrom)改版,加入了多种reid,原来ZQPei提供的针对行人跟踪的权重ckpt.t7不能直接使用。以下记录如何在新版中使用osnet reid模型,以及使用ZQPei ckpt.t7模型的方法。
经验证,新版Yolov5_DeepSort_Pytorch,用osnet_x1_0, osnet_ain_x1_0均可运行,性能和ZQPei模型差不多,但速度慢。大约40ms:20ms/帧的差别。可能的原因,osnet 匹配图像大,256x128(h,w), ZQPei ckpt图像小128x64(h,w)。
尝试将ZQPei模型写成新版reid方式,在导入模型权重时,提示丢弃了两个不匹配的层,结果运行速度偏慢,40ms,性能也不及原来的好。
mikel-brostrom引入KaiyangZhou提供的reid,其使用方法如下。
如何导入torchreid:
将KaiyangZhou github克隆下来,放到Yolov5_DeepSort_Pytorch/deep_sort/deep目录下,目录名为改reid,即Yolov5_DeepSort_Pytorch/deep_sort/deep/reid。
假定已经安装了conda和虚拟环境,且安装好Yolov5_DeepSort_Pytorch所需的模块,进入reid目录,运行
python setup.py develop
如此即安装好torchreid,可以在程序中加入import torchreid。
从KaiyangZhou的github中,Model zoo里下载权重文件,如osnet_x1_0.pth,放到checkpoint目录:Yolov5_DeepSort_Pytorch/deep_sort/deep/checkpoint。
修改deep_sort.yaml
DEEPSORT:
MODEL_TYPE: "osnet_x1_0"
REID_CKPT: '~/Yolov5_DeepSort_Pytorch/deep_sort/deep/checkpoint/osnet_x1_0_imagenet.pth'
MAX_DIST: 0.1 # 0.2 The matching threshold. Samples with larger distance are considered an invalid match
MAX_IOU_DISTANCE: 0.7 # 0.7 Gating threshold. Associations with cost larger than this value are disregarded.
MAX_AGE: 90 # 30 Maximum number of missed misses before a track is deleted
N_INIT: 3 # 3 Number of frames that a track remains in initialization phase
NN_BUDGET: 100 # 100 Maximum size of the appearance descriptors gallery
MIN_CONFIDENCE: 0.75
NMS_MAX_OVERLAP: 1.0
track.py中指定reid模型
parser.add_argument('--deep_sort_model', type=str, default='osnet_x1_0')
在deep_sort.yaml文件中给出权重文件路径,可跳过从网上下载权重的过程,直接从本地下载。如此可运行osnet reid之deepsort跟踪程序track.py。
将ZQPei模型添加到reid中的办法:
修改model.py,在py文件中添加
def mymodel(num_classes=751, pretrained=True, loss='softmax', **kwargs):
model = Net(
num_classes=num_classes,
pretrained = pretrained,
loss = 'softmax',
**kwargs
)
return model
为避免名称发生冲突,将原来的model.py改成mymodel.py。
在deep_sort/deep/reid/torchreid/models/__init__.py
中添加:
from .mymodel import *
在
__model_factory = {
下面添加自己的模型:
'mymodel': mymodel
在deep_sort/deep/reid/torchreid/utils/feature_extractor.py中添加
from deep_sort.deep.reid.torchreid.models.mymodel import Net
__init__函数中修改image_size:
image_size=(128, 64),
好像可以运行了。但这样导入的权重似乎有些不妥,运行速度较慢,性能也变差一点。
这样修改feature_extractor.py似乎更好些,能保持ckpt.t7之能力:
from __future__ import absolute_import
import numpy as np
import torch
import torchvision.transforms as T
from PIL import Image
import cv2
from torchreid.utils import (
check_isfile, load_pretrained_weights, compute_model_complexity
)
from torchreid.models import build_model
import logging
from deep_sort.deep.reid.torchreid.models.mymodel import Net
class FeatureExtractor(object):
def __init__(
self,
model_name='',
model_path='',
image_size=(64, 128), # 256, 128 w, h
pixel_mean=[0.485, 0.456, 0.406],
pixel_std=[0.229, 0.224, 0.225],
pixel_norm=True,
device='cuda'
):
self.net = Net(pretrained =True)
if model_path and check_isfile(model_path):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
state_dict = torch.load(model_path, map_location=torch.device(self.device))['net_dict']
self.net.load_state_dict(state_dict)
self.net.eval()
self.size = (64,128) # self.size = (64,128) , self.size(width, height)
import torchvision.transforms as transforms
self.norm = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
device = torch.device(device)
self.net.to(device)
def _preprocess(self, im_crops):
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)
return features # .cpu().numpy()