0
点赞
收藏
分享

微信扫一扫

新版Yolov5_DeepSort_Pytorch使用ZQPei行人模型的方法

由于号称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()
举报

相关推荐

YOLOV5+deepsort 行人检测与跟踪

0 条评论