目前基于Person_reid_baseline_pytorch的教程更多的是关于如何run起来这个工程,但如果想应用到其他细分项目上的教程比较少,数据集较为单一,且命令和格式复杂,所以本文将详细描述如何制作自建的数据集并选择不同的backbone进行训练(因为ResNet-50难以应用于落地)。
一、自建数据集的分类
1、训练集train+val,测试集query+gallery
以train为例,train文件夹下是各个类别图片存放的文件夹:
其中train+val是所有的训练图片,是全部图片的一部分,而另一部分图片则组成了测试集gallery+query,所以训练集和测试集加起来才是所有类别文件夹。val中各类别下是训练图片中每个分类下选取的一张图片,query中是测试图片中每个分类下,每个摄像机选取一张图片,但由于都是直播截图,所以只在每个类别文件夹下放置了一张图片,而其他的测试图片都放在gallery中。
2、命名规则
随便打开一张图片,其命名为
(1)0004代表其标签编号,属于第四个类
(2)c1代表这张图片来自第一个camera
(3)s6代表这张图片来自这个摄像头的非连续的第六个视频片段(自制数据集无区分)
(4)016996表示是c1s6的第16996帧
(5)02表示这是该帧中使用DPM得到的第2个检测框,若是手工标注的则为00
注意事项:即使是自制数据集没有那么多摄像头,摄像头数量c那里,也要给他随机random.randint(1,6)几个值,否则在test生成参数的时候,会得不到正确结果,反应如下
3、数据集制作
(1)首先根据不同视频中的帧图片进行上半身检测,将得到的位置信息在图片中进行截取,获得每个视频中的上半身图片,当同个视频中存在不同的人物时,需要一一分成不同的文件夹。
(2)将每个文件夹重命名为0001…0199等,每个文件夹中的图片进行重命名为上述命名规则。
这里放一个批量更改各类文件夹下的图片的脚本:
# -*- coding: utf-8 -*-
import os
import random
# 设定文件路径
path1 = './20230306/'
i = 3000
# 对目录下的文件进行遍历
for category in os.listdir(path1):
# 设置新文件夹名
folder_name = str(i)
folder_path = os.path.join(path1, folder_name)
# if not os.path.exists(folder_path):
# 重命名
Newdir = os.path.join(path1,folder_name.zfill(4))
os.rename(os.path.join(path1, category), Newdir)
#os.rename(folder_path,os.path.join(folder_path.zfill(4)))
file_path = Newdir#os.path.join(path1,category)
j = 1
# 对目录下的文件进行遍历
for file in os.listdir(file_path):
# print(file)
# 判断是否是文件
# if os.path.isfile(os.path.join(file_path, file)):
# 设置新文件名
file_name = str(folder_name) + '_' + 'c'+str(random.randint(1,6))+'s1_' + str(j).zfill(6) + '.jpg'
# 重命名
new_file_name = os.path.join(file_path, file_name)
os.rename(os.path.join(file_path, file), new_file_name)
print(new_file_name)
j += 1
i += 1
# 结束
print("End")
(3)按比例将全部文件夹分成训练集(train+val)和测试集(query+gallery)
这里也附一个脚本:
import cv2
import os
import glob
import shutil
import random
#path = r"D:\xxx\32\*.jpg"
def select_some_folder(filepath,target_filepath,filepath1,target_filepath1):
path_dir = os.listdir(filepath)
pick_number = 100 #这里的100就是你想从所有数据中(未抽取前所有数据都在train和val里)随机抽多少class进测试集(query和gallery)
sample = random.sample(path_dir,pick_number)
for name in sample:
shutil.move(filepath+name,target_filepath+name)
shutil.move(filepath1+name,target_filepath1+name)
if __name__ =="__main__":
fileDir = "./train/"
tarDir = "./val/"
query_dir = "./query/"
gallery_dir = "./gallery/"
select_some_folder(fileDir,gallery_dir,tarDir,query_dir)
(4)val的每个文件夹与train同名,从train中取一张图片置入;query的每个文件夹与gallery同名,从gallery中取一张图片置入。
这里同样还附一个脚本:
import cv2
import os
import glob
import shutil
import random
#path = r"D:xxx\32\*.jpg"
def select_one_image(filepath,target_filepath):
path_dir = os.listdir(filepath)
pick_number = 1
sample = random.sample(path_dir,pick_number)
for name in sample:
shutil.move(filepath+name,target_filepath+name)
if __name__ =="__main__":
fileDir = "./train/"
tarDir = "./val/"
for oneDir in os.listdir(fileDir):
onefileDir = fileDir+oneDir+"/"
onetarDir = tarDir+oneDir+"/"
if not os.path.exists(onetarDir):
os.mkdir(onetarDir)
select_one_image(onefileDir,onetarDir)
二、训练
python train.py
--data_dir 数据集所在位置
--name 训练结果所在文件夹名称
--gpu_ids 使用哪几个gpu
在model.py中第75行
原文中说需要将class_num=751进行修改,这里的751使用的是market-1501数据集中的分类,但其实数据集名称做好的情况下,不需修改。
内置torch.utils.data.Dataloader将创建迭代器dataloader[‘train’]自动计算数据集中的类别数量。
训练结束后将获得一个.pth模型文件,存储了训练得到的特征数据,后续test将使用这个文件。
三、测试
Python test.py
--name 注意这个name代表选择哪个训练结果,和train中的含义稍许不同
--which_epoch 选择模型文件中的哪个epoch
--test_dir 数据集path
(1)若训练正常的话,最后的torch.size[a,b]中的a是你测试集的类别数,如果不对的话需要进行检查哪里出错。
(2)如上图所示,如果rank后的值都为0的话,则是数据命名时的camera只设置了1,需要取更多camera值。
会生成一个result.mat文件,其中包含了query和gallery各图片的匹配信息,以字典的形式检索。
四、Demo.py
Python demo.py
--query_index 从第几个query队列中选择图片进行特征对比
--test_dir 数据集Path
运行后会从result.mat中调用特征信息,最后会生成一个show.png,第一张为你query_index给定的检索图片,后面10张是检索到的最匹配的图片,上标绿色的为同一类,上标红色的为不同类。