百度西交大大数据菁英班目标检测竞赛

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

数据介绍

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

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

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

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

总体思路

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

!pip install paddlex


import 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 os

f = 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 shuffle

f = 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往期赛+新赛资料包

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