fanacio / yolox-master Goto Github PK

import os
import random

images_path = "JPEGImages/"
xmls_path = "Annotations/"
train_val_txt_path = "ImageSets/Main/"
val_percent = 0.1

images_list = os.listdir(images_path)
random.shuffle(images_list)

#　划分训练集和验证集的数量
train_images_count = int((1-val_percent)*len(images_list))
val_images_count = int(val_percent*len(images_list))

#　生成训练集的train.txt文件
train_txt = open(os.path.join(train_val_txt_path,"train.txt"),"w")
train_count = 0
for i in range(train_images_count):
    text = images_list[i].split(".jpg")[0] + "\n"
    train_txt.write(text)
    train_count+=1
    print("train_count: " + str(train_count))
train_txt.close()

#　生成验证集的val.txt文件
val_txt = open(os.path.join(train_val_txt_path,"val.txt"),"w")
val_count = 0
for i in range(val_images_count):
    text = images_list[train_images_count + i].split(".jpg")[0] + "\n"
    val_txt.write(text)
    val_count+=1
    print("val_count: " + str(val_count))
val_txt.close()

        #修改标签信息，类别后都要加一个逗号
        VOC_CLASSES = (
            "head",
        )

        class Exp(MyExp):
            def __init__(self):
                super(Exp, self).__init__()
                self.num_classes = 1           #修改类别数量，voc原始为20
                self.depth = 0.33
                self.width = 0.50
                self.warmup_epochs = 1

        class Exp(BaseExp):
            def __init__(self):
                super().__init__()

                # ---------------- model config ---------------- #
                # detect classes number of model
                self.num_classes = 1            #修改类别数量，COCO原始为80

        with wait_for_the_master(local_rank):
                    dataset = VOCDetection(
                        data_dir="/home/dataset/head_det/",   #训练集所在的绝对路径（docker内）
                        image_sets=[('train')],               #训练集名称
                        img_size=self.input_size,
                        preproc=TrainTransform(
                            max_labels=100,       #表示单图中最多出现的目标个数，修改为100，原始值为50
                            flip_prob=self.flip_prob,

        #原始的rootpath中带有年份信息，自己标注的数据集中没有年份信息，所以删掉:
                for name in image_sets:
                    rootpath = self.root
                    for line in open(
                        os.path.join(rootpath, "ImageSets", "Main", name + ".txt")
                    ):
                        self.ids.append((rootpath, line.strip()))

        def get_eval_loader(self, batch_size, is_distributed, testdev=False, legacy=False):
            from yolox.data import VOCDetection, ValTransform

            valdataset = VOCDetection(
                data_dir="/home/dataset/head_det/",        #验证集的绝对路径（docker中）
                image_sets=[('val')],                      #验证集的名字
                img_size=self.test_size,
                preproc=ValTransform(legacy=legacy),
            )

        class Exp(MyExp):
            def __init__(self):
                super(Exp, self).__init__()
                self.num_classes = 1           #修改类别数量，voc原始为20
                self.depth = 0.33
                self.width = 0.50
                self.warmup_epochs = 1

        self.input_size = (640, 640)  # (height, width)

        self.test_size = (640, 640)

        class Exp(BaseExp):
            def __init__(self):
                super().__init__()

                # ---------------- model config ---------------- #
                # detect classes number of model
                self.num_classes = 1            #修改类别数量，COCO原始为80
                # factor of model depth
                self.depth = 0.33               #修改深度，原始值为1.00
                # factor of model width
                self.width = 0.5               #修改宽度，原始值为1.00

        self.max_epoch = 300       #总的迭代的epoch
        # minimum learning rate during warmup
        self.warmup_lr = 0
        self.min_lr_ratio = 0.05
        # learning rate for one image. During training, lr will multiply batchsize.
        self.basic_lr_per_img = 0.01 / 64.0
        # name of LRScheduler
        self.scheduler = "yoloxwarmcos"
        # last #epoch to close augmention like mosaic
        self.no_aug_epochs = 15
        # apply EMA during training
        self.ema = True

        # weight decay of optimizer
        self.weight_decay = 5e-4
        # momentum of optimizer
        self.momentum = 0.9
        # log period in iter, for example,
        # if set to 1, user could see log every iteration.
        self.print_interval = 5        #表示每迭代5次保存一下pth权重，原始值为10
        # eval period in epoch, for example,
        # if set to 1, model will be evaluate after every epoch.
        self.eval_interval = 5        #表示每迭代5次对验证集做一次验证，原始值为10
        # save history checkpoint or not.

        def _do_python_eval(self, output_dir="output", iou=0.5):
                rootpath = self.root                #去掉了rootpath中原有的年份信息
                name = self.image_set[0]            #由于自有数据集中没有年份信息，所以名字不是image_set[0][1]，而是image_set[0]
                annopath = os.path.join(rootpath, "Annotations", "{:s}.xml")
                imagesetfile = os.path.join(rootpath, "ImageSets", "Main", name + ".txt")
                cachedir = os.path.join(
                    self.root, "annotations_cache"              #去掉cachedir中原有的年份信息
                )
                if not os.path.exists(cachedir):
                    os.makedirs(cachedir)
                aps = []
                # The PASCAL VOC metric changed in 2010
                use_07_metric = True  #改为True，之前是根据年份信息进行判断的
                print("Eval IoU : {:.2f}".format(iou))

        def _get_voc_results_file_template(self):
                filename = "comp4_det_test" + "_{:s}.txt"
                filedir = os.path.join(self.root, "results")        #由于自建数据集中没有年份信息，所以将年份信息的参数删除
                if not os.path.exists(filedir):