【yolact】训练数据集

无敌牛战士

已于 2024-01-08 20:22:44 修改

阅读量788

点赞数 2

文章标签： python 开发语言

于 2023-10-20 17:32:19 首次发布

本文链接：https://blog.csdn.net/m0_71106777/article/details/133949472

版权

一、下载yolact代码

yolact在github中的源码地址：GitHub - dbolya/yolact: A simple, fully convolutional model for real-time instance segmentation.

下载权重库

二、准备自己的数据集

1、准备好用labelme标注好的图片和.json格式文件

2、将数据集转化为coco数据集

在文件夹下创建labelme2coco.py文件，代码来源github链接

import os
import json
import numpy as np
import glob
import shutil
from sklearn.model_selection import train_test_split
np.random.seed(41)

#0为背景，此处根据你数据集的类别来修改key
classname_to_id = {"1": 1}

class Lableme2CoCo:

    def __init__(self):
        self.images = []
        self.annotations = []
        self.categories = []
        self.img_id = 0
        self.ann_id = 0

    def save_coco_json(self, instance, save_path):
        json.dump(instance, open(save_path, 'w', encoding='utf-8'), ensure_ascii=False, indent=1)  # indent=2 更加美观显示

    # 由json文件构建COCO
    def to_coco(self, json_path_list):
        self._init_categories()
        for json_path in json_path_list:
            obj = self.read_jsonfile(json_path)
            self.images.append(self._image(obj, json_path))
            shapes = obj['shapes']
            for shape in shapes:
                annotation = self._annotation(shape)
                self.annotations.append(annotation)
                self.ann_id += 1
            self.img_id += 1
        instance = {}
        instance['info'] = 'spytensor created'
        instance['license'] = ['license']
        instance['images'] = self.images
        instance['annotations'] = self.annotations
        instance['categories'] = self.categories
        return instance

    # 构建类别
    def _init_categories(self):
        for k, v in classname_to_id.items():
            category = {}
            category['id'] = v
            category['name'] = k
            self.categories.append(category)

    # 构建COCO的image字段
    def _image(self, obj, path):
        image = {}
        from labelme import utils
        img_x = utils.img_b64_to_arr(obj['imageData'])
        h, w = img_x.shape[:-1]
        image['height'] = h
        image['width'] = w
        image['id'] = self.img_id
        image['file_name'] = os.path.basename(path).replace(".json", ".jpg")
        return image

    # 构建COCO的annotation字段
    def _annotation(self, shape):
        label = shape['label']
        points = shape['points']
        annotation = {}
        annotation['id'] = self.ann_id
        annotation['image_id'] = self.img_id
        annotation['category_id'] = int(classname_to_id[label])
        annotation['segmentation'] = [np.asarray(points).flatten().tolist()]
        annotation['bbox'] = self._get_box(points)
        annotation['iscrowd'] = 0
        annotation['area'] = 1.0
        return annotation

    # 读取json文件，返回一个json对象
    def read_jsonfile(self, path):
        with open(path, "r", encoding='utf-8') as f:
            return json.load(f)

    # COCO的格式： [x1,y1,w,h] 对应COCO的bbox格式
    def _get_box(self, points):
        min_x = min_y = np.inf
        max_x = max_y = 0
        for x, y in points:
            min_x = min(min_x, x)
            min_y = min(min_y, y)
            max_x = max(max_x, x)
            max_y = max(max_y, y)
        return [min_x, min_y, max_x - min_x, max_y - min_y]


if __name__ == '__main__':
    labelme_path = "labelme/"   # 此处根据你的数据集地址来修改
    saved_coco_path = "./"
    # 创建文件
    if not os.path.exists("%scoco/annotations/"%saved_coco_path):
        os.makedirs("%scoco/annotations/"%saved_coco_path)
    if not os.path.exists("%scoco/images/train2017/"%saved_coco_path):
        os.makedirs("%scoco/images/train2017"%saved_coco_path)
    if not os.path.exists("%scoco/images/val2017/"%saved_coco_path):
        os.makedirs("%scoco/images/val2017"%saved_coco_path)
    # 获取images目录下所有的joson文件列表
    json_list_path = glob.glob(labelme_path + "/*.json")
    # 数据划分,这里没有区分val2017和tran2017目录，所有图片都放在images目录下
    train_path, val_path = train_test_split(json_list_path, test_size=0.12)
    print("train_n:", len(train_path), 'val_n:', len(val_path))

    # 把训练集转化为COCO的json格式
    l2c_train = Lableme2CoCo()
    train_instance = l2c_train.to_coco(train_path)
    l2c_train.save_coco_json(train_instance, '%scoco/annotations/instances_train2017.json'%saved_coco_path)
    for file in train_path:
        shutil.copy(file.replace("json","jpg"),"%scoco/images/train2017/"%saved_coco_path)
    for file in val_path:
        shutil.copy(file.replace("json","jpg"),"%scoco/images/val2017/"%saved_coco_path)

    # 把验证集转化为COCO的json格式
    l2c_val = Lableme2CoCo()
    val_instance = l2c_val.to_coco(val_path)
    l2c_val.save_coco_json(val_instance, '%scoco/annotations/instances_val2017.json'%saved_coco_path)

这个代码中需要修改两个地方

（1）按照自己的标签顺序修改标签

#0为背景，此处根据你数据集的类别来修改key
classname_to_id = {"1": 1}

（2）根据自己数据集的保存路径和要保存的位置更改路径

if __name__ == '__main__':
    labelme_path = "labelme/"   # 此处根据你的数据集地址来修改
    saved_coco_path = "./"

结果输出产生两个文件

至此数据集准备完成

三、修改config.py文件

1、修改标签名字和顺序

其中类名中COCO_CLASSES =('hym',)最后结尾的“，”不可缺少

2、修改图片或视频的储存位置

3、修改coco_base_config

'num_classes'：数量需要包含background，即=label数+1

此个max_iter可以控制训练轮数

四、开始训练

train.py参数

--resume：接着中断的权重进行训练

–batch_size：大小取决于你的显卡，报错就调小一点
–num_workers：线程数，一般都是需要设置为 “0”
–lr：学习率
–save_folder：训练完后模型保存的路径
–config：config配置文件选择，none为默认配置
–save_interval：每10000次迭代保存模型权重
–validation_epoch：每2个epoch进行一次验证计算
–dataset：数据集路径

parser.add_argument('--batch_size', default=4, type=int,  ########################################################################
                    help='Batch size for training')

parser.add_argument('--num_workers', default=0, type=int,  #######################################################################
                    help='Number of workers used in dataloading')

parser.add_argument('--lr', '--learning_rate', default=None, type=float,
                    help='Initial learning rate. Leave as None to read this from the config.')
parser.add_argument('--momentum', default=None, type=float,
                    help='Momentum for SGD. Leave as None to read this from the config.')
parser.add_argument('--decay', '--weight_decay', default=None, type=float,
                    help='Weight decay for SGD. Leave as None to read this from the config.')
                    
parser.add_argument('--save_folder', default='weights/',
                    help='Directory for saving checkpoint models.')

parser.add_argument('--config', default=None,    
                    help='The config object to use.')
                    
parser.add_argument('--save_interval', default=10000, type=int,
                    help='The number of iterations between saving the model.')
                    
parser.add_argument('--validation_size', default=5000, type=int,
                    help='The number of images to use for validation.')
                    
parser.add_argument('--validation_epoch', default=2, type=int,
                    help='Output validation information every n iterations. If -1, do no validation.')

parser.add_argument('--dataset', default=None, type=str,
                    help='If specified, override the dataset specified in the config with this one (example: coco2017_dataset).')

执行train.py代码

python train.py --config=yolact_base_config

开始后等T收敛了就可以按Ctrl+c终止代码保存模型

五、测试

eval.py参数

–trained_model：选择模型文件
–top_k：保存置信率最高的前15个目标
–config：选择config配置文件
–image：对单张图像进行预测，路径为单张图像
–images：对多张图像/图像路径文件夹进行预测
–video：指定视频进行预测
–score_threshold ：剔除掉低于0.15置信度的目标
–dataset：数据集路径

    parser.add_argument('--trained_model',##########################################################################################
                        default='weights/yolact_base_290_16000.pth', type=str,
                        help='Trained state_dict file path to open. If "interrupt", this will open the interrupt file.')

    parser.add_argument('--top_k', default=15, type=int, #################################################################
                        help='Further restrict the number of predictions to parse')
                        
    parser.add_argument('--config', default=None,   ##############################################################################
                        help='The config object to use.')

    parser.add_argument('--image', default=None, type=str,
                        help='A path to an image to use for display.')

    parser.add_argument('--images', default='D:/y/XM_TEST/yolact-master/data/coco/real_test', type=str,################################################################################
                        help='An input folder of images and output folder to save detected images. Should be in the format input->output.')
                        
    parser.add_argument('--video', default=None, type=str,
                        help='A path to a video to evaluate on. Passing in a number will use that index webcam.')
                        
    parser.add_argument('--score_threshold', default=0.15, type=float,   #################################################################################
                        help='Detections with a score under this threshold will not be considered. This currently only works in display mode.')
                        
    parser.add_argument('--dataset', default=None, type=str,
                        help='If specified, override the dataset specified in the config with this one (example: coco2017_dataset).')

执行eval.py代码

python eval.py --trained_model=weights/yolact_base_169_169_interrupt.pth --score_threshold=0.15 --top_k=15 --image=1.jpg

预测结果

六、可能遇到的问题

1、ZeroDivisionError: division by zero

该错误为train.py文件中的epoch_size为0

解决办法

（1）修改代码为

epoch_size = len(dataset)+1 // args.batch_size

但是此方法会导致数据不不精确

（2）增大训练数据集的数量，我最开始只用了8张图片所以才导致此错误

2、UnicodeDecodeError: 'gbk' codec can't decode byte 0xb4 in position 8: illegal multibyte sequence

该错误是由于数据集中的保存路径含有中文或者label标签为中文，我直接将两个都改为了不含中文的，就解决了

3、RuntimeError:Expected a ‘cuda‘ device type for generator but found ‘cpu‘

解决办法：在train.py中

将：

data_loader = data.DataLoader(dataset, args.batch_size,
num_workers=args.num_workers,
shuffle=True, collate_fn=detection_collate,
pin_memory=True)

改为：

data_loader = data.DataLoader(dataset, args.batch_size,
num_workers=args.num_workers,
shuffle=True, collate_fn=detection_collate,
pin_memory=True,generator=torch.Generator(device = 'cuda'))

data_loader = data.DataLoader(dataset, args.batch_size, # 创建一个数据加载器，将dataset分成批次加载
num_workers=args.num_workers, # 设置加载数据的并行工作进程数
shuffle=True, # 打乱数据的顺序
collate_fn=detection_collate, # 将样本列表转换为批次的函数
pin_memory=True, # 将数据加载到 CUDA 固定内存中
generator=torch.Generator(device=‘cuda’)) # 设置生成器的设备为 CUDA，用于随机打乱数据的顺序