YOLOv5-7.0可以用来做实例分割的任务了！！！用完感觉实在是666啊
目录
项目介绍
数据标注及处理
json转换txt
切分训练集、测试集、验证集
修改配置文件
模型训练和推理
后处理
项目介绍
本文章主要目的有两个：
- 用yolov5分割网络训练自己的数据
- 处理yolov5的分割结果，将分割的图像裁剪出来
我的项目是需要识别图一里面这些小块，将每个小块裁剪出来，旋转成水平角度后再进行下一步的操作。因项目保密原因，就用模糊的图片代替，见谅见谅。下面展示了效果图，如果你的项目需要实现的功能跟我类似，可参考参考
图一
图二
图三
图四
图片说明：
- 图一是原图
- 图二是yolov5检测后的图片，它用其他颜色将目标的mask给标出来，并画出目标的最外矩形框
- 图三（后处理）是计算目标四个点坐标，并在原图上画出来
- 图四（后处理）是将目标旋转摆正并切分成小图
数据标注及处理
- 标注工具：labelme
- 标注文件：json格式
- 训练数据要求：坐标归一化的txt文件
从作者提供的样例数据coco128-seg（下载链接：https://ultralytics.com/assets/coco128-seg.zip），可以看到txt文件的内容，分别是类别下标，归一化的坐标，中间用空格分割，不同目标物体用换行符
json转换txt
如何将我们用labelme标注的json文件转化为对应的格式呢？
import json
import os
import argparse
from tqdm import tqdm
def 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 shutil
import random
import os
import 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
修改配置文件
1、data文件夹里面有yaml文件，下面图片是data/coco128-seg.yaml的内容。
path：是上面--save-dir切分图片存放的路径；
train、val、test分别对于images里面的文件夹，按实际填入；
names：是类别名称和赋予的下标，跟上面转txt顺序相同
2、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文件，每一次训练都会生成对应的权重文件
2、模型推理
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个子元素，每个元素都是多个坐标点构成，坐标点是目标预测出来的轮廓坐标值
后处理需要做的步骤有：
- 坐标反归一化：segments的坐标和txt的格式相同，都是归一化的坐标值，需要转换为图片真实的坐标值
- 获取四个点：多个坐标点计算出左上、右上、左下、右下的点，并将他们按顺时针的顺序输出
- 旋转摆正：已知四个点，可计算得出角度，将目标摆正后保存成小图
代码献出
# 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.shape
for 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 cv2
import numpy as np
def 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_img
def 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
撒花完结！！！