在做一个简单的项目时，先训练了一个安全帽检测模型，后面有训练了一个安全衣检测模型，后来又说要将两个模型，合成一个模型，其实说起来也很简单，就是重新标注一下数据，清理一下，然后再训练一下就行，但是实际操作其来却然人头皮发麻。
首先第一步就是数据的重新标注，想象一下，安全帽数据集有七八千张图片，安全衣数据集有三四千张图片，这些图片合在一起就一万多张。重新标注！！！！！！！！！！！！！
下面说一下，实际操作过程中的部分问题及解决方法，和部分代码展示。
第一步，就是数据集构建阶段，当时最开始搞的时候，就想着偷懒，直接将两个数据集合在一起，改一下标签序号后直接训练，但是训练结果就是，基本上就是有头没衣服，有衣服没头这种。
更改txt标签文件，每行的类别序号，代码如下。
import os
import linecache
YOLOV5_LABEL_ROOT = r"E:\code\detect\yolov5\dataset\all_safety\all_safety\3\\" # yolov5 导出的推理图片的 txt，需要改变的路径
DATASET_LABEL_ROOT = r"E:\code\detect\yolov5\dataset\all_safety\all_safety\4\\" # 数据集的保存路径
if __name__ == '__main__':
yolo_file = os.listdir(YOLOV5_LABEL_ROOT)
# 遍历文件里面有 .txt 结尾的
for file_name in yolo_file:
# 判断 txt 文件才进行读取
if not file_name.endswith(".txt"):
continue
file_path = YOLOV5_LABEL_ROOT + file_name
with open(file_path, "r+") as f:
for line in f.readlines():
# line = linecache.getline(file_path, 1)
line1 = line[:1]
line2 = line[1:]
# line3 = line1.replace("1", "3") # 0为原来的内容，3为要替换的内容
line3 = line1.replace("4", "0")
line = line3 + line2
data_path = DATASET_LABEL_ROOT + file_name
print(data_path)
# 汇总到数据集的标注文件
with open(data_path, "a") as fd:
fd.write(line)第二种，就是开始给每张图签进行打标，然后先用安全帽的模型标注安全衣数据集中的安全帽，然后用安全衣的模型标注安全帽数据集的安全衣。这句话有点绕，仔细看一下。
这种方法，需要在更改的时候仔细点，改标签的时候别改错，反复检查几遍就行。
还有一点就是，数据清理，由于安全帽数据集更大一些，所以训练时的权重会更偏向安全帽一些，导致安全衣的训练效果极差，所以需要去除一些安全帽或者没带安全帽的相关数据集（这个前提是安全帽数据集中的负样本，涵盖的安全衣及其他衣服少，均是比较密级的头的情况，所以此方式要根据实际情况而定）
第三种方法就是全部重新打标，为了减少工作量，使用一下代码直接生成标签，在对数据集进行微调，即可满足需求。代码如下（同时加载两个模型实现推理生成标签）
import os
from pathlib import Path
import torch
from models.common import DetectMultiBackend
from utils.dataloaders import LoadImages
from utils.general import LOGGER, Profile, check_img_size, colorstr, cv2, non_max_suppression, scale_boxes, xyxy2xywh
from utils.plots import Annotator, colors
from utils.torch_utils import select_device
def detect(weights_hat, weights_safety, source, save_dir, half, line_thickness):
if not os.path.exists(save_dir):
os.mkdir(save_dir)
device = select_device('0') # 获取设备
# =============================================== 读取是否穿戴安全帽模型 ===================================================
models = DetectMultiBackend(weights_hat, device=device, fp16=half)
stride, names, pt = models.stride, models.names, models.pt
print('pt', pt)
# =============================================== 读取是否穿安全衣模型 ====================================================
models_safety = DetectMultiBackend(weights_safety, device=device, fp16=half)
names_safety, pt_safety = models_safety.names, models_safety.pt
imgsz = check_img_size(imgsz=640, s=stride) # check image size
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt)
# dataset_safety = LoadImages(source, img_size=imgsz, stride=stride, auto=pt_safety)
seen, windows, dt = 0, [], (Profile())
# 正式推理 从上面的LoadImages()可以看到每次只输入单张图片
# 处理每一张图片或者视频的格式
for path, im, im0s, vid_cap, s in dataset:
# path 是图片或视频的路径
# im 是进行resize + pad之后的图片
# im0s 是原尺寸图片
# vid_cap 当读取的图片为None时，读取视频为视频源
# torch.from_numpy()的作用是将生成的数组转换为张量。
im = torch.from_numpy(im).to(models.device)
# 图片设置为FP16/FP32
im = im.half() if models.fp16 else im.float()
im /= 255 # 归一化
# 如果没有batch_size的话则在最前面添加一个轴
if len(im.shape) == 3:
# 增加一个维度
im = im[None]
with dt: # 使用with以及dt是为了方便后续计算推理时间
pred = models(im)
pred_safety = models_safety(im)
# 调用utils/general.py文件中的non_max_suppression类进行非极大值抑制
pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False)
pred_safety = non_max_suppression(pred_safety, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False)
# 后续保存或者打印预测信息
# 对每张图片进行处理 将pred映射回原图
for i, det in enumerate(pred):
seen += 1
p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
p = Path(p) # 当前路径（yolov5/data/images）
save_path = str(save_dir + p.name) # 图片或视频的保存路径（如runs/detect/exp/***.jpg）
# 设置保存框的坐标的txt文件路径，每张图片对应一个框坐标信息
txt_path = str(save_dir + 'labels' + p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')
# 设置打印图片的信息
s += '%gx%g ' % im.shape[2:]
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
annotator = Annotator(im0, line_width=line_thickness, example=str(names))
if len(det):
det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round() #将预测信息映射到原图
# 打印检测到的类别数量
for c in det[:, 5].unique():
n = (det[:, 5] == c).sum()
s += f"{n} {names[int(c)]}{'s' * (n > 1)},"
# name = ["0", "dog", '1', '3']
# 保存结果： txt/图片画框/crop-image
for *xyxy, conf, cls in reversed(det):
# 将每个图片的预测信息分别存入save_dir/labels下的xxx.txt中 每行: class_id + score + xywh
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(
-1).tolist() # normalized xywh
if int(cls) == 2:
line = (cls - 1, *xywh) # label format 这里的类别序号需要注意一下
else:
line = (cls, *xywh) # label format 这里的类别序号需要注意一下
with open(f'{txt_path}.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
# 在原图中画框
for *xyxy, conf, cls in reversed(det):
c = int(cls) # 整数类
print(names)
label = (f'{names[c]} {conf:.2f}')
annotator.box_label(xyxy, label, color=colors(c, True))
im0 = annotator.result()
# 保存图片
if dataset.mode == 'image':
cv2.imwrite(save_path, im0)
# =====================================加载安全衣检测模型对图片进行推理======================================================
# 后续保存或者打印预测信息
# 对每张图片进行处理 将pred映射回原图
for i, det in enumerate(pred_safety):
seen += 1
p, im0, frame = path, im0.copy(), getattr(dataset, 'frame', 0)
p = Path(p) # 当前路径（yolov5/data/images）
save_path = str(save_dir + p.name) # 图片或视频的保存路径（如runs/detect/exp/***.jpg）
# 设置保存框的坐标的txt文件路径，每张图片对应一个框坐标信息
txt_path = str(save_dir + 'labels' + p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')
# 设置打印图片的信息
s += '%gx%g ' % im.shape[2:]
annotator = Annotator(im0, line_width=line_thickness, example=str(names_safety))
if len(det):
det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round() # 将预测信息映射到原图
# 打印检测到的类别数量
for c in det[:, 5].unique():
n = (det[:, 5] == c).sum()
s += f"{n} {names_safety[int(c)]}{'s' * (n > 1)},"
# 保存结果： txt/图片画框/crop-image
for *xyxy, conf, cls in reversed(det):
# 将每个图片的预测信息分别存入save_dir/labels下的xxx.txt中 每行: class_id + score + xywh
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(
-1).tolist() # normalized xywh
line = (cls+2, *xywh) # label format 这里类别序号也需要注意一下
with open(f'{txt_path}.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
# 在原图中画框
for *xyxy, conf, cls in reversed(det):
c = int(cls) # 整数类
print(names_safety)
label = (f'{names_safety[c]} {conf:.2f}')
annotator.box_label(xyxy, label, color=colors(c, True))
im0_safety = annotator.result()
# 保存图片
if dataset.mode == 'image':
cv2.imwrite(save_path, im0_safety)
# 打印推理时间，图片大小，以及检测到的类别
LOGGER.info(f"{s}{''if len(det) else '(no detection),'}{dt.dt * 1E3:.1f}ms")
# 打印输出图像的保存路劲
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")
if __name__ == "__main__":
with torch.no_grad():
half = False # 是否使用半精度推理（F16），使用半精度可以提高检测速度
line_thickness = 3 # 画框线条的粗细
weights_hat = r'E:\code\detect\yolov5\runs\train\best_hat\weights/best_hat.pt/'
weights_safety = r'E:\code\detect\yolov5\runs\train\best_clothes\weights/best_clothes.pt/'
# weights1 = r'E:\code\detect\yolov5\yolov5s.pt/'
source = r'E:\code\detect\yolov5\data\images/'
source = str(source)
# labels = ['hat', 'other_clothes', 'person', 'reflective_clothes']
save_dir = r'E:\code\detect\yolov5\runs\detect\a/'
detect(weights_hat, weights_safety, source, save_dir, half, line_thickness)画框结果如下：（）
标签如下：
