飞桨打比赛：极市计算机视觉开发者榜单大赛——安全帽识别

import osimport jsonimport randomimport argparseimport paddlex as pdxfrom paddlex import transforms as Tdef split_dataset(data_dir, save_dir, split_num, label_list):
    '''
    分割数据集

    参数
        data_dir 数据目录
        save_dir 保存目录
        split_num 验证集数量
        label_list 标签列表

    返回
        train_file 训练集列表文件
        val_file 验证集列表文件
        label_file 标签列表文件
    '''

    # 遍历数据文件
    jpg_files = []
    xml_files = []
    data_files = os.listdir(data_dir)    for data_file in data_files:        if data_file.endswith('.jpg'):
            jpg_files.append(data_file)        elif data_file.endswith('.xml'):
            xml_files.append(data_file)    # 排序并组成数据对
    datas = []
    jpg_files.sort()
    xml_files.sort()    for jpg, xml in zip(jpg_files, xml_files):
        datas.append(f'{jpg} {xml}\n')    # 打乱顺序
    random.shuffle(datas)    # 写入文件
    train_file = os.path.join(save_dir, 'train.txt')
    val_file = os.path.join(save_dir, 'val.txt')
    label_file = os.path.join(save_dir, 'label_list.txt')    with open(train_file, 'w', encoding='UTF-8') as f:        for data in datas[:-split_num]:
            f.write(data)    with open(val_file, 'w', encoding='UTF-8') as f:        for data in datas[-split_num:]:
            f.write(data)    with open(label_file, 'w', encoding='UTF-8') as f:        for label in label_list:
            f.write(f'{label}\n')    # 输出信息
    states = {}
    states['label_list'] = label_list
    states['datas_num'] = len(datas)
    states['train_num'] = len(datas[:-split_num])
    states['val_num'] = len(datas[-split_num:])
    states['train_file'] = train_file
    states['val_file'] = val_file
    states['label_flie'] = label_file    print(json.dumps(states, indent=4))    return train_file, val_file, label_fileif __name__ == '__main__':    # 命令行参数
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_dir', '-d', default='/home/data/831', type=str)
    parser.add_argument('--label_list', '-l', default='person,head,hat')
    parser.add_argument('--save_dir', '-s', default='/project/train/src_repo', type=str)
    parser.add_argument('--ckpt_dir', '-c', default='/project/train/models', type=str)
    parser.add_argument('--split_num', '-n', default=50, type=int)
    args = parser.parse_known_args()[0]    # 打印命令行参数
    print(json.dumps(vars(args), indent=4))    # 参数转换
    data_dir = args.data_dir
    save_dir = args.save_dir
    ckpt_dir = args.ckpt_dir
    split_num = args.split_num
    label_list = args.label_list.split(',')    
    # 切分数据集
    train_file, val_file, label_file = split_dataset(data_dir, save_dir, split_num, label_list)    # 训练集数据增强
    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=[320, 352, 384, 416, 448, 480, 512, 544, 576, 608],
            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=480, interp='CUBIC'),
        T.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ])    # 训练集
    train_dataset = pdx.datasets.VOCDetection(
        data_dir=data_dir,
        file_list=train_file,
        label_list=label_file,
        transforms=train_transforms,
        num_workers=0,
        shuffle=True
    )    # 验证集
    eval_dataset = pdx.datasets.VOCDetection(
        data_dir=data_dir,
        file_list=val_file,
        label_list=label_file,
        transforms=eval_transforms,
        num_workers=0,
        shuffle=False
    )    # 检测模型
    model = pdx.det.PicoDet(
        num_classes=len(train_dataset.labels),
        backbone='ESNet_s',
        nms_score_threshold=.025,
        nms_topk=1000,
        nms_keep_topk=100,
        nms_iou_threshold=.6
    )    # 模型训练
    model.train(
        num_epochs=20,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        train_batch_size=16,
        pretrain_weights='COCO',
        learning_rate=0.0001,
        warmup_steps=200,
        warmup_start_lr=0.0,
        save_interval_epochs=3,
        lr_decay_epochs=[15, 18],
        save_dir=ckpt_dir,
        use_vdl=False
    )

$ python /project/train/src_repo/train.py -n 100

import jsonimport numpy as npimport paddlex as pdx# 模型路径model_path = '/project/train/models/best_model'# 阈值threshold = 0.5def init():
    '''
    初始化

    返回
        model PaddleX 模型
    '''
    model = pdx.load_model(model_path)    return modeldef process_image(handle=None, input_image=None, args=None, **kwargs):
    '''
    处理图像

    参数
        handle init 函数的返回值
        input_image 输入图像 (CHW / BGR)
        args / **kwargs 其他参数
    
    返回
        result json 格式的结果
    '''
    results = handle.predict(input_image)
    objects = []    for dt in np.array(results):
        cname, bbox, score = dt['category'], dt['bbox'], dt['score']        if score > threshold:
            objects.append({                "x": round(bbox[0]),                "y": round(bbox[1]),                "width": round(bbox[2]),                "height": round(bbox[3]),                "confidence": score,                "name": cname
            })
    result = {        "model_data": {            "objects": objects
        }
    }    return json.dumps(result, indent=4)if __name__ == '__main__':    import cv2    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--img', '-i', default='/home/data/831/helmet_38270.jpg', type=str)
    args = parser.parse_known_args()[0]
    model = init()
    img = cv2.imread(args.img)
    process_image(model, img)