来源：投稿 作者：LSC
编辑：学姐

数据介绍

数据集共包括40000张训练图像和1000张测试图像，每张训练图像对应xml标注文件：

共包含3类：0:'head', 1:'helmet', 2:'person'。

提交格式要求，提交名为pred_result.txt的文件——每一行代表一个目标，每一行内容分别表示：图像名 置信度 xmin ymin xmax ymax类别

「限制只能使用paddle框架和aistudio平台运行代码」

总体思路

使用paddlex框架，模型选取ppyolov2模型。

!pip install paddleximport paddlex as pdx
from paddlex import transforms as T

## 数据增强train_transforms = T.Compose([T.MixupImage(mixup_epoch=-1), T.RandomDistort(),T.RandomExpand(im_padding_value=[123.675, 116.28, 103.53]), T.RandomCrop(),T.RandomHorizontalFlip(), T.BatchRandomResize(target_sizes=[192, 224, 256, 288, 320, 352, 384, 416, 448, 480, 512],interp='RANDOM'), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])eval_transforms = T.Compose([T.Resize(target_size=320, interp='CUBIC'), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])import osf = open("work/total.txt", "w", encoding="utf-8")
for i in os.listdir("work/helmet/train/images/"):voc = "annotations/" + i[:-3] + "xml"  f.write("images/" + i + "\t" + voc + "\n")
f.close()# 最后一行是错误格式，手动删除

f = open("work/test.txt", "w", encoding="utf-8")
for i in os.listdir("work/helmet/test/images/"):voc = "annotations/" + i[:-3] + "xml"  f.write("images/" + i + "\t" + voc + "\n")
f.close()from sklearn.utils import shufflef = open("work/total.txt", "r", encoding="utf-8")
total = f.readlines()ratio = 0.9
total = shuffle(total, random_state = 100)
train_len = int(len(total) * ratio)train = total[:train_len]
val = total[train_len:]f1 = open("work/train.txt", "w", encoding="utf-8")
for i in train:f1.write(i)
f1.close()f2 = open("work/val.txt", "w", encoding="utf-8")
for i in val:f2.write(i)
f2.close()f.close()

#手动创建label.txt

数据导入

train_dataset = pdx.datasets.VOCDetection(data_dir='work/helmet/train/',file_list='work/train.txt',label_list='work/label.txt',transforms=train_transforms,shuffle=True)test_dataset = pdx.datasets.VOCDetection(data_dir='work/helmet/test/',file_list='work/test.txt',label_list='work/label.txt',transforms=eval_transforms)eval_dataset = pdx.datasets.VOCDetection(data_dir='work/helmet/train/',file_list='work/val.txt',label_list='work/label.txt',transforms=eval_transforms)

# 在训练集上聚类生成9个anchor
anchors = train_dataset.cluster_yolo_anchor(num_anchors=9, image_size=608)
anchor_masks = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]

#开始训练
num_classes = len(train_dataset.labels)
model = pdx.det.PPYOLOv2(num_classes=num_classes,backbone='ResNet101_vd_dcn',anchors=anchors,anchor_masks=anchor_masks,label_smooth=True)model.train(num_epochs=100,train_dataset=train_dataset,train_batch_size=8,eval_dataset=eval_dataset,pretrain_weights='COCO',learning_rate=0.005 / 12,warmup_steps=500,warmup_start_lr=0.0,save_interval_epochs=5,# lr_decay_epochs=[25, 75],save_dir='output1/',use_vdl=False,early_stop=True,
early_stop_patience=5)

# 使用之前最好的模型继续训练
model.train(num_epochs=100,train_dataset=train_dataset,train_batch_size=8,eval_dataset=eval_dataset,# pretrain_weights='COCO',learning_rate=0.005 / 12,warmup_steps=500,warmup_start_lr=0.0,save_interval_epochs=5,# lr_decay_epochs=[25, 75],save_dir='output2/',pretrain_weights='output1/best_model/model.pdparams',use_vdl=False,early_stop=True,
early_stop_patience=5)

# 导入最好的模型，评估模型效果
model = pdx.load_model("output1/best_model")
model.evaluate(eval_dataset, batch_size=8, metric=None, return_details=False)

# 模型推理，生成的两个文本文件就是最终提交的结果
image_dirs = 'work/helmet/test/images/'
f1 = open("work/pred_result1.txt", "w", encoding="utf-8") # 只写阈值大于0.5的
f2 = open("work/pred_result2.txt", "w", encoding="utf-8") # 全部写 
for image_name in os.listdir(image_dirs):result = model.predict(image_dirs + image_name)for i in range(len(result)):xmin, ymin = int(result[i]['bbox'][0]), int(result[i]['bbox'][1])xmax, ymax = int(xmin + result[i]['bbox'][2]), int(ymin + result[i]['bbox'][3])if result[i]['score'] >= 0.5:f1.write(image_name[:-4] + " " + str(result[i]['score']) + " " + str(xmin) + " " + str(ymin) + " " + str(xmax) + " " + str(ymax) \+ " " + str(result[i]['category_id']) + "\n")f2.write(image_name[:-4] + " " + str(result[i]['score']) + " " + str(xmin) + " " + str(ymin) + " " + str(xmax) + " " + str(ymax) \+ " " + str(result[i]['category_id']) + "\n")
f1.close()
f2.close()

最终mAP值达到62.77648。

后续可以使用PaddleDetection框架进行优化，选取其中的ppyoloplus模型或者PaddleYOLO框架中的yolov5、yolov6、yolox、yolov7模型。ppyoloplus模型优化后的效果可以达到65%以上。

关注下方《学姐带你玩AI》🚀🚀🚀

回复“all in”免费领取kaggle往期赛+新赛资料包

码字不易，欢迎大家点赞评论收藏！