YOLOv5-7.0可以用来做实例分割的任务了！！！用完感觉实在是666啊

目录

项目介绍

 数据标注及处理

        json转换txt 

        切分训练集、测试集、验证集

修改配置文件

模型训练和推理

后处理

项目介绍

本文章主要目的有两个：

1. 用yolov5分割网络训练自己的数据
2. 处理yolov5的分割结果，将分割的图像裁剪出来

我的项目是需要识别图一里面这些小块，将每个小块裁剪出来，旋转成水平角度后再进行下一步的操作。因项目保密原因，就用模糊的图片代替，见谅见谅。下面展示了效果图，如果你的项目需要实现的功能跟我类似，可参考参考

 图一

 图二

 图三

图四

图片说明：

• 图一是原图
• 图二是yolov5检测后的图片，它用其他颜色将目标的mask给标出来，并画出目标的最外矩形框
• 图三（后处理）是计算目标四个点坐标，并在原图上画出来
• 图四（后处理）是将目标旋转摆正并切分成小图

 数据标注及处理

1. 标注工具：labelme
2. 标注文件：JSON格式
3. 训练数据要求：坐标归一化的txt文件

 从作者提供的样例数据coco128-seg（下载链接：https://ultralytics.com/assets/coco128-seg.zip），可以看到txt文件的内容，分别是类别下标，归一化的坐标，中间用空格分割，不同目标物体用换行符

        json转换txt 

如何将我们用labelme标注的json文件转化为对应的格式呢？

import jsonimport osimport argparsefrom tqdm import tqdmdef convert_label_json(json_dir, save_dir, classes):    json_paths = os.listdir(json_dir)    classes = classes.split(',')    for json_path in tqdm(json_paths):    # for json_path in json_paths:        path = os.path.join(json_dir,json_path)        with open(path,'r') as load_f:            json_dict = json.load(load_f)        h, w = json_dict['imageHeight'], json_dict['imageWidth']        # save txt path        txt_path = os.path.join(save_dir, json_path.replace('json', 'txt'))        txt_file = open(txt_path, 'w')        for shape_dict in json_dict['shapes']:            label = shape_dict['label']            label_index = classes.index(label)            points = shape_dict['points']            points_nor_list = []            for point in points:                points_nor_list.append(point[0]/w)                points_nor_list.append(point[1]/h)            points_nor_list = list(map(lambda x:str(x),points_nor_list))            points_nor_str = ' '.join(points_nor_list)                        label_str = str(label_index) + ' ' +points_nor_str + '\n'            txt_file.writelines(label_str)if __name__ == "__main__":    """    python json2txt_nomalize.py --json-dir my_datasets/color_rings/jsons --save-dir my_datasets/color_rings/txts --classes "cat,dogs"    """    parser = argparse.ArgumentParser(description='json convert to txt params')    parser.add_argument('--json-dir', type=str, help='json path dir')    parser.add_argument('--save-dir', type=str, help='txt save dir')    parser.add_argument('--classes', type=str, help='classes')    args = parser.parse_args()    json_dir = args.json_dir    save_dir = args.save_dir    classes = args.classes    convert_label_json(json_dir, save_dir, classes)

脚本说明：

    --json-dir：标注的纯json目录；

    --save-dir：要保存的txt文件目录；

    --classes：类别名称，它的类别顺序跟后面的配置文件顺序相同，如类别cat,dog，执行命令可以这么写

Python json2txt_nomalize.py --json-dir my_datasets/color_rings/jsons --save-dir my_datasets/color_rings/txts --classes "cat,dog"

        切分训练集、测试集、验证集

# 将图片和标注数据按比例切分为 训练集和测试集import shutilimport randomimport osimport argparse# 检查文件夹是否存在def mkdir(path):    if not os.path.exists(path):        os.makedirs(path)def main(image_dir, txt_dir, save_dir):    # 创建文件夹    mkdir(save_dir)    images_dir = os.path.join(save_dir, 'images')    labels_dir = os.path.join(save_dir, 'labels')    img_train_path = os.path.join(images_dir, 'train')    img_test_path = os.path.join(images_dir, 'test')    img_val_path = os.path.join(images_dir, 'val')    label_train_path = os.path.join(labels_dir, 'train')    label_test_path = os.path.join(labels_dir, 'test')    label_val_path = os.path.join(labels_dir, 'val')    mkdir(images_dir);mkdir(labels_dir);mkdir(img_train_path);mkdir(img_test_path);mkdir(img_val_path);mkdir(label_train_path);mkdir(label_test_path);mkdir(label_val_path);    # 数据集划分比例，训练集75%，验证集15%，测试集15%，按需修改    train_percent = 0.8    val_percent = 0.1    test_percent = 0.1    total_txt = os.listdir(txt_dir)    num_txt = len(total_txt)    list_all_txt = range(num_txt)  # 范围 range(0, num)    num_train = int(num_txt * train_percent)    num_val = int(num_txt * val_percent)    num_test = num_txt - num_train - num_val    train = random.sample(list_all_txt, num_train)    # 在全部数据集中取出train    val_test = [i for i in list_all_txt if not i in train]    # 再从val_test取出num_val个元素，val_test剩下的元素就是test    val = random.sample(val_test, num_val)    print("训练集数目：{}, 验证集数目：{},测试集数目：{}".fORMat(len(train), len(val), len(val_test) - len(val)))    for i in list_all_txt:        name = total_txt[i][:-4]        srcImage = os.path.join(image_dir, name+'.jpg')        srcLabel = os.path.join(txt_dir, name + '.txt')        if i in train:            dst_train_Image = os.path.join(img_train_path, name + '.jpg')            dst_train_Label = os.path.join(label_train_path, name + '.txt')            shutil.copyfile(srcImage, dst_train_Image)            shutil.copyfile(srcLabel, dst_train_Label)        elif i in val:            dst_val_Image = os.path.join(img_val_path, name + '.jpg')            dst_val_Label = os.path.join(label_val_path, name + '.txt')            shutil.copyfile(srcImage, dst_val_Image)            shutil.copyfile(srcLabel, dst_val_Label)        else:            dst_test_Image = os.path.join(img_test_path, name + '.jpg')            dst_test_Label = os.path.join(label_test_path, name + '.txt')            shutil.copyfile(srcImage, dst_test_Image)            shutil.copyfile(srcLabel, dst_test_Label)if __name__ == '__main__':    """    python split_datasets.py --image-dir my_datasets/color_rings/imgs --txt-dir my_datasets/color_rings/txts --save-dir my_datasets/color_rings/train_data    """    parser = argparse.ArgumentParser(description='split datasets to train,val,test params')    parser.add_argument('--image-dir', type=str, help='image path dir')    parser.add_argument('--txt-dir', type=str, help='txt path dir')    parser.add_argument('--save-dir', type=str, help='save dir')    args = parser.parse_args()    image_dir = args.image_dir    txt_dir = args.txt_dir    save_dir = args.save_dir    main(image_dir, txt_dir, save_dir)

脚本说明：

    --image-dir：训练图片目录；

    --txt-dir：上一步生成txt的目录； 

    --save-dir：切分数据集存放路径，执行命令样例：

python split_datasets.py --image-dir my_datasets/color_rings/imgs --txt-dir my_datasets/color_rings/txts --save-dir my_datasets/color_rings/train_data

执行后在存放路径可以看到，自动生成images和labels两个文件夹，两个文件夹里面有三个文件夹：train\test\val

修改配置文件

data文件夹里面有yaml文件，下面图片是data/coco128-seg.yaml的内容。

    path：是上面--save-dir切分图片存放的路径；

    train、val、test分别对于images里面的文件夹，按实际填入；

    names：是类别名称和赋予的下标，跟上面转txt顺序相同

models/segment文件夹也有yaml文件，如果你使用yolov5m模型，就修改yolov5m-seg.yaml文件的nc，如果有两个类别，nc就修改成2

模型训练和推理

1、训练执行命令

python segment/train.py --epochs 300 --data coco128-seg.yaml --weights yolov5m-seg.pt --img 640 --cfg models/segment/yolov5m-seg.yaml --batch-size 16 --device 2

执行命令说明：指明配置文件、预训练权重路径等，具体参数查看train.py文件

结果：在runs目录生成train-seg文件，每一次训练都会生成对应的权重文件

模型推理

python segment/predict.py --weight ./runs/train-seg/exp2/weights/best.pt --source ./my_datasets/color_rings/train_data/images/test/000030.jpg

执行命令说明：指明权重路径和预测的图片或者文件夹，具体参数查看predict.py文件

结果：在runs目录生成predict-seg目录，保存了上面图二的结果图

后处理

重要的后处理来了！！！

segment/predict.py，约169行附近，将预测坐标保存在txt文件。打印segments的维度，他是一个list，如果预测的图片中有6个目标，那么list包含了6个子元素，每个元素都是多个坐标点构成，坐标点是目标预测出来的轮廓坐标值

后处理需要做的步骤有：

1. 坐标反归一化：segments的坐标和txt的格式相同，都是归一化的坐标值，需要转换为图片真实的坐标值
2. 获取四个点：多个坐标点计算出左上、右上、左下、右下的点，并将他们按顺时针的顺序输出
3. 旋转摆正：已知四个点，可计算得出角度，将目标摆正后保存成小图 

代码献出

# segments是分割的坐标点segments = [    scale_segments(im0.shape if retina_masks else im.shape[2:], x, im0.shape, normalize=True)    for x in reversed(masks2segments(masks))]new_segments = []  # 用来装反归一化后的坐标image_list = []    # 切割的小图im0_h, im0_w, im0_c = im0.shapefor k, seg_list in enumerate(segments):    # 将归一化的点转换为坐标点    new_seg_list = []    for s_point in seg_list:        pt1, pt2 = s_point        new_pt1 = int(pt1 * im0_w)        new_pt2 = int(pt2 * im0_h)        new_seg_list.append([new_pt1, new_pt2])    rect = cv2.minAreaRect(np.array(new_seg_list))  # 得到最小外接矩形的（中心(x,y), (宽,高), 旋转角度）    seg_bbox = cv2.boxPoints(rect)  # 获取最小外接矩形的4个顶点坐标(ps: cv2.boxPoints(rect) for OpenCV 3.x)    seg_bbox = np.int0(seg_bbox)    if np.linalg.norm(seg_bbox[0] - seg_bbox[1]) < 5 or np.linalg.norm(seg_bbox[3] - seg_bbox[0]) < 5:        continue    # 坐标点排序    box1 = sorted(seg_bbox, key=lambda x: (x[1], x[0]))    # 将坐标点按照顺时针方向来排序，box的从左往右从上到下排序    if box1[0][0] > box1[1][0]:        box1[0], box1[1] = box1[1], box1[0]    if box1[2][0] < box1[3][0]:        box1[2], box1[3] = box1[3], box1[2]    if box1[0][1] > box1[1][1]:        box1[0], box1[1], box1[2], box1[3] = box1[1], box1[2], box1[3], box1[0]    box1_list = [b.tolist() for b in box1] # 坐标转换为list格式    new_segments.append(box1_list)    tmp_box = copy.deepcopy(np.array(box1)).astype(np.float32)    partImg_array = image_crop_tools.get_rotate_crop_image(im0, tmp_box)    image_list.append(partImg_array)    # cv2.imwrite(str(k)+'.jpg', partImg_array)  # 保存小图    # 在原图上画出分割图像# src_image = im0.copy()# for ns_box in new_segments:#     cv2.drawContours(src_image, [np.array(ns_box)], -1, (0, 255, 0), 2)# cv2.imwrite('1.jpg', src_image)

 代码说明：该部分脚本复制在segment/predict.py文件，可以放在if save_txt的同一级别下面。其中注释#保存小图，是保存文章开头图四的图片。注释#在原图上画出分割图像，是文章开头图三的图像。

旋转部分用到了image_crop_tools.get_rotate_crop_image函数，主要用来做角度计算和图片摆正，代码如下：

import cv2import numpy as npdef get_rotate_crop_image(img, points):    """    根据坐标点截取图像    :param img:     :param points:     :return:     """    h, w, _ = img.shape    left = int(np.min(points[:, 0]))    right = int(np.max(points[:, 0]))    top = int(np.min(points[:, 1]))    bottom = int(np.max(points[:, 1]))    img_crop = img[top:bottom, left:right, :].copy()    points[:, 0] = points[:, 0] - left    points[:, 1] = points[:, 1] - top    img_crop_width = int(np.linalg.norm(points[0] - points[1]))    img_crop_height = int(np.linalg.norm(points[0] - points[3]))    pts_std = np.float32([[0, 0], [img_crop_width, 0], [img_crop_width, img_crop_height], [0, img_crop_height]])    M = cv2.getPerspectiveTransform(points, pts_std)    dst_img = cv2.warpPerspective(        img_crop,        M, (img_crop_width, img_crop_height),        borderMode=cv2.BORDER_REPLICATE)    dst_img_height, dst_img_width = dst_img.shape[0:2]    if dst_img_height * 1.0 / dst_img_width >= 1:#         pass#         print(dst_img_height * 1.0 / dst_img_width,dst_img_height,dst_img_width,'*-'*10)        dst_img = np.rot90(dst_img,-1)   #-1为逆时针，1为顺时针。        return dst_imgdef sorted_boxes(dt_boxes):    """    坐标点排序    """       num_boxes = dt_boxes.shape[0]    sorted_boxes = sorted(dt_boxes, key=lambda x: (x[0][1], x[0][0]))    _boxes = list(sorted_boxes)    for i in range(num_boxes - 1):        if abs(_boxes[i+1][0][1] - _boxes[i][0][1]) < 10 and \            (_boxes[i + 1][0][0] < _boxes[i][0][0]):            tmp = _boxes[i]            _boxes[i] = _boxes[i + 1]            _boxes[i + 1] = tmp       return _boxes

 撒花完结！！！

来源地址：https://blog.csdn.net/jin__9981/article/details/128385498