小白也能学会的跨摄像头实现目标追踪和个体识别（YOLOV5+ReID）

目录

0.项目流程

1.训练ReID模型

2.训练YOLOV5模型

3.在YOLOV5中添加ReID模型

4.项目运行步骤

（1）下载ReID模型

（2）创建虚拟环境

（3）在虚拟环境下载pytorch

（4）下载所有需要的依赖

（5）修改search中的训练参数

（6）跳转到当前目录/yolov5_reid/model

（7）运行结果

---

0.项目流程

项目地址：

https://github.com/yaofeino1/YOLOV5-ReID

在目标追踪与个体识别任务中，我们可以结合 ReID 模型与 YOLO 模型来实现跨摄像头的精准识别。具体流程如下：首先，利用训练好的 YOLO 模型对视频进行逐帧检测，快速定位出感兴趣的目标区域，比如人物。然后，借助已经训练好的 ReID 模型，提取这些由 YOLO 模型检测到的目标区域的特征向量，这些特征向量就像是每个人物的独特 “数字指纹”。

当我们需要在视频中搜索特定目标时，先对待搜索目标图片使用 ReID 模型提取其特征向量，再将该特征向量与视频中检测到的目标区域的特征向量进行匹配。如果在视频中检测到与待搜索目标特征向量匹配的人物，就对其进行打标签和画框操作，从而实现目标追踪。

基于这一原理，如果我们搜索多个人，即输入多个人物图片并提取各自的特征向量，在视频检测过程中，通过与视频中检测到的目标区域特征向量一一匹配，就可以直接在视频里找到所有人对应的 ID，实现个体识别。这样一来，无论是单目标追踪还是多目标的个体识别，都能高效完成。

1.训练ReID模型

训练ReID模型可以参考这个博客

手把手教你如何从零开始行人重识别（ReID的训练和测试）-CSDN博客

其中训练好的模型是保存在model/ft_ResNet50这个文件夹里的。

2.训练YOLOV5模型

该项目是基于https://github.com/ultralytics/yolov5/tree/v5.0修改的

该项目使用的yolov5（v5.0）版本进行修改的，如果有同学想要自己训练目标识别模型，最好是用这个github仓库里面的代码进行训练。

3.在YOLOV5中添加ReID模型

我将ReID模型封装成一个类，这个类的作用是实现特征提取，保存在extracter.py里，下面是代码

import argparse
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.backends.cudnn as cudnn
import numpy as np
import torchvision
from torchvision import transforms
import math
from model import ft_net, ft_net_dense, ft_net_hr, ft_net_swin, ft_net_swinv2, ft_net_efficient, ft_net_NAS, ft_net_convnext, PCB, PCB_test
from utils import fuse_all_conv_bn
import yaml
import os

class FeatureExtractor:
    def __init__(self, config_path, which_epoch='last', gpu_ids='0', ms='1'):
        # 解析配置文件
        with open(config_path, 'r') as stream:
            config = yaml.load(stream, Loader=yaml.FullLoader)

        self.opt = argparse.Namespace()
        for key, value in config.items():
            setattr(self.opt, key, value)
        self.opt.which_epoch = which_epoch
        self.opt.gpu_ids = gpu_ids
        self.opt.ms = ms

        str_ids = self.opt.gpu_ids.split(',')
        self.gpu_ids = []
        for str_id in str_ids:
            id = int(str_id)
            if id >= 0:
                self.gpu_ids.append(id)

        str_ms = self.opt.ms.split(',')
        self.ms = []
        for s in str_ms:
            s_f = float(s)
            self.ms.append(math.sqrt(s_f))

        if len(self.gpu_ids) > 0:
            torch.cuda.set_device(self.gpu_ids[0])
            cudnn.benchmark = True

        # 选择模型结构
        if self.opt.use_dense:
            model_structure = ft_net_dense(self.opt.nclasses, stride=self.opt.stride, linear_num=self.opt.linear_num)
        elif self.opt.use_NAS:
            model_structure = ft_net_NAS(self.opt.nclasses, linear_num=self.opt.linear_num)
        elif self.opt.use_swin:
            model_structure = ft_net_swin(self.opt.nclasses, linear_num=self.opt.linear_num)
        elif self.opt.use_swinv2:
            model_structure = ft_net_swinv2(self.opt.nclasses, (self.opt.h, self.opt.w), linear_num=self.opt.linear_num)
        elif self.opt.use_convnext:
            model_structure = ft_net_convnext(self.opt.nclasses, linear_num=self.opt.linear_num)
        elif self.opt.use_efficient:
            model_structure = ft_net_efficient(self.opt.nclasses, linear_num=self.opt.linear_num)
        elif self.opt.use_hr:
            model_structure = ft_net_hr(self.opt.nclasses, linear_num=self.opt.linear_num)
        else:
            model_structure = ft_net(self.opt.nclasses, stride=self.opt.stride, ibn=self.opt.ibn,
                                     linear_num=self.opt.linear_num)

        if self.opt.PCB:
            model_structure = PCB(self.opt.nclasses)

        self.model = self.load_network(model_structure)

        # Remove the final fc layer and classifier layer
        if self.opt.PCB:
            self.model = PCB_test(self.model)
        else:
            self.model.classifier.classifier = nn.Sequential()

        # Change to test mode
        self.model = self.model.eval()
        if torch.cuda.is_available():
            self.model = self.model.cuda()

        self.model = fuse_all_conv_bn(self.model)

        # 定义图像预处理转换
        if self.opt.use_swin:
            self.h, self.w = 224, 224
        else:
            self.h, self.w = 256, 128

        self.data_transforms = transforms.Compose([
            transforms.Resize((self.h, self.w), interpolation=3),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

        if self.opt.PCB:
            self.data_transforms = transforms.Compose([
                transforms.Resize((384, 192), interpolation=3),
                transforms.ToTensor(