from os import getcwd
from sklearn.model_selection import train_test_split
import json
import glob
wd = getcwd()
"labelme标注的json 数据集转为pytorch版yolov4的训练集"
classes = ["aircraft","oiltank"]

image_ids = glob.glob(r"LabelmeData/*jpg")
print(image_ids)
train_list_file = open('data/train.txt', 'w')
val_list_file = open('data/val.txt', 'w')
def convert_annotation(image_id, list_file):
    jsonfile=open('%s.json' % (image_id))
    in_file = json.load(jsonfile)

    for i in range(0,len(in_file["shapes"])):
        object=in_file["shapes"][i]
        cls=object["label"]
        points=object["points"]
        xmin=int(points[0][0])
        ymin=int(points[0][1])
        xmax=int(points[1][0])
        ymax=int(points[1][1])
        if cls not in classes:
            print("cls not in classes")
            continue
        cls_id = classes.index(cls)
        b = (xmin, ymin, xmax, ymax)
        list_file.write(" " + ",".join([str(a) for a in b]) + ',' + str(cls_id))
    jsonfile.close()

def ChangeData2TXT(image_List,dataFile):
    for image_id in image_List:
        dataFile.write('%s' % (image_id.split('\\')[-1]))
        convert_annotation(image_id.split('.')[0], dataFile)
        dataFile.write('\n')
    dataFile.close()
trainval_files, test_files = train_test_split(image_ids, test_size=0.2, random_state=55)
ChangeData2TXT(trainval_files,train_list_file)
ChangeData2TXT(test_files,val_list_file)

安装运行需要的包

参照requirements.txt安装本机没有的包，版本不一定要保持一致，只要后期不报错就没有事。

修改类别

将coco.names和voc.names里面的类别修改为自己数据集的类别（默认是coco.names，都改了肯定没有错。），顺序和labelme2txt.py中的classes顺序保持一致。

修改配置文件cfg.py

Cfg.use_darknet_cfg = False

Cfg.batch = 2(根据自己的显卡修改，我的显卡是8G的最多可以训练里2个batch)。

Cfg.subdivisions = 1

修改models.py

将51行和53行的inplace=True改为inplace=False。如果不修改，训练的时候会报个错误。

修改train.py文件

找到526行，这个方法的参数是对cfg.py里面参数的更新。

主要修改的参数如下：

parser.add_argument('-g', '--gpu', metavar='G', type=str, default='0',
help='GPU', dest='gpu')#设置GPU使用的GPU

parser.add_argument('-dir', '--data-dir', type=str, default="LabelmeData",
help='dataset dir', dest='dataset_dir')#图片所在的文件夹。

parser.add_argument('-pretrained', type=str, default="data/yolov4.conv.137.pth", help='pretrained yolov4.conv.137')#设置预训练权重文件的路径。

parser.add_argument('-classes', type=int, default=80, help='dataset classes')#物体类别数。

parser.add_argument('-train_label_path', dest='train_label', type=str, default='data/train.txt', help="train label path")#训练集存放的路径。

注释415行到440行的代码，这段代码在验证的时候一直报错，我找不到原因。后续找到原因再更新。

将以上的内容修改完成后就可以点击run开始训练了。

测试

测试主要修改models.py的代码。将下面的代码从449行替换。


if __name__ == "__main__":
 
    import sys
 
    import cv2
 
 
 
    namesfile = None
 
    n_classes=2
 
    weightfile="checkpoints/Yolov4_epoch151.pth" 
 
    imgfile="data/aircraft_4.jpg"#待测试的图片
 
    width=608
 
    height=608
 
    model = Yolov4(yolov4conv137weight=None, n_classes=n_classes, inference=True)
 
 
 
    pretrained_dict = torch.load(weightfile, map_location=torch.device('cpu'))
 
#如果使用GPU则改为：
 
#pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda'))

    model.load_state_dict(pretrained_dict)

    use_cuda = True
    if use_cuda:
        model.cuda()

    img = cv2.imread(imgfile)

    # Inference input size is 416*416 does not mean training size is the same
    # Training size could be 608*608 or even other sizes
    # Optional inference sizes:
    #   Hight in {320, 416, 512, 608, ... 320 + 96 * n}
    #   Width in {320, 416, 512, 608, ... 320 + 96 * m}
    sized = cv2.resize(img, (width, height))
    sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)

    from tool.utils import load_class_names, plot_boxes_cv2
    from tool.torch_utils import do_detect

    for i in range(2): # This 'for' loop is for speed check
                        # Because the first iteration is usually longer
        boxes = do_detect(model, sized, 0.4, 0.6, use_cuda)

    if namesfile == None:
        if n_classes == 2:
            namesfile = 'data/voc.names'
        elif n_classes == 80:
            namesfile = 'data/coco.names'
        else:
            print("please give namefile")

    class_names = load_class_names(namesfile)
    resultImg=plot_boxes_cv2(img, boxes[0], 'predictions.jpg', class_names)
    cv2.imshow("image",resultImg)
    cv2.waitKey(0);

测试结果：

源码链接：https://download.csdn.net/download/hhhhhhhhhhwwwwwwwwww/12821500

参考文章：

深入浅出Yolo系列之Yolov3&Yolov4核心基础知识完整讲解

https://blog.csdn.net/nan355655600/article/details/106246625?utm_medium=distribute.pc_relevant.none-task-blog-BlogCommendFromMachineLearnPai2-2.nonecase&depth_1-utm_source=distribute.pc_relevant.none-task-blog-BlogCommendFromMachineLearnPai2-2.nonecase

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/菜鸟追梦旅行/article/detail/128662