YOLOv7-7u-seg实现矿井下轨道和人图像分割实例
一、代码获取
首先进入YOLOv7下载图像分割版本，并且在releases中获取YOLOv5s-seg或者YOLOv7-seg，将这两个模型文件加到YOLOv7根目录中，然后去网上下载Arial.ttf文件，因为运行detect.py进行代码测试时会去网上下载这些文件，不如自己先下载。
二、构建数据集
新建conda环境anylabelme，并且进入该虚拟环境。
conda create -n labelme python=3.8
conda activate labelme
在虚拟环境中安装该软件，安装成功输入anylabelme打开软件。
pip install labelme
labelme
进行标注，记得点击自动保存。
标注完成后，将json格式标签转换为YOLO格式，也就是txt结尾。
import base64
import random
import shutil
from tqdm import tqdm
import math
import json
import os
import numpy as np
import PIL.Image
import PIL.ImageDraw
import cv2
class ConvertManager(object):
def __init__(self):
pass
def base64_to_numpy(self, img_bs64):
img_bs64 = base64.b64decode(img_bs64)
img_array = np.frombuffer(img_bs64, np.uint8)
cv2_img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
return cv2_img
@classmethod
def load_labels(cls, name_file):
'''
load names from file.one name one line
:param name_file:
:return:
'''
with open(name_file, 'r') as f:
lines = f.read().rstrip('\n').split('\n')
return lines
def get_class_names_from_all_json(self, json_dir):
classnames = []
for file in os.listdir(json_dir):
if not file.endswith('.json'):
continue
with open(os.path.join(json_dir, file), 'r', encoding='utf-8') as f:
data_dict = json.load(f)
for shape in data_dict['shapes']:
if not shape['label'] in classnames:
classnames.append(shape['label'])
return classnames
def create_save_dir(self, save_dir):
images_dir = os.path.join(save_dir, 'images')
labels_dir = os.path.join(save_dir, 'labels')
if not os.path.exists(save_dir):
os.makedirs(save_dir)
os.mkdir(images_dir)
os.mkdir(labels_dir)
else:
if not os.path.exists(images_dir):
os.mkdir(images_dir)
if not os.path.exists(labels_dir):
os.mkdir(labels_dir)
return images_dir + os.sep, labels_dir + os.sep
def save_list(self, data_list, save_file):
with open(save_file, 'w') as f:
f.write('\n'.join(data_list))
def __rectangle_points_to_polygon(self, points):
xmin = 0
ymin = 0
xmax = 0
ymax = 0
if points[0][0] > points[1][0]:
xmax = points[0][0]
ymax = points[0][1]
xmin = points[1][0]
ymin = points[1][1]
else:
xmax = points[1][0]
ymax = points[1][1]
xmin = points[0][0]
ymin = points[0][1]
return [[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]]
def convert_dataset(self, json_dir, json_list, images_dir, labels_dir, names, save_mode='train'):
images_dir = os.path.join(images_dir, save_mode)+os.sep
labels_dir = os.path.join(labels_dir, save_mode)+os.sep
if not os.path.exists(images_dir):
os.mkdir(images_dir)
if not os.path.exists(labels_dir):
os.mkdir(labels_dir)
for file in tqdm(json_list):
with open(os.path.join(json_dir, file), 'r', encoding='utf-8') as f:
data_dict = json.load(f)
image_file = os.path.join(json_dir, os.path.basename(data_dict['imagePath']))
if os.path.exists(image_file):
shutil.copyfile(image_file, images_dir + os.path.basename(image_file))
else:
imageData = data_dict.get('imageData')
if not imageData:
imageData = base64.b64encode(imageData).decode('utf-8')
img = self.img_b64_to_arr(imageData)
PIL.Image.fromarray(img).save(images_dir + file[:-4] + 'png')
# convert to txt
width = data_dict['imageWidth']
height = data_dict['imageHeight']
line_list = []
for shape in data_dict['shapes']:
data_list = []
data_list.append(str(names.index(shape['label'])))
if shape['shape_type'] == 'rectangle':
points = self.__rectangle_points_to_polygon(shape['points'])
for point in points:
data_list.append(str(point[0] / width))
data_list.append(str(point[1] / height))
elif shape['shape_type'] == 'polygon':
points = shape['points']
for point in points:
data_list.append(str(point[0] / width))
data_list.append(str(point[1] / height))
line_list.append(' '.join(data_list))
self.save_list(line_list, labels_dir + file[:-4] + "txt")
def split_train_val_test_dataset(self, file_list, train_ratio=0.9, trainval_ratio=0.9, need_test_dataset=False,
shuffle_list=True):
if shuffle_list:
random.shuffle(file_list)
total_file_count = len(file_list)
train_list = []
val_list = []
test_list = []
if need_test_dataset:
trainval_count = int(total_file_count * trainval_ratio)
trainval_list = file_list[:trainval_count]
test_list = file_list[trainval_count:]
train_count = int(train_ratio * len(trainval_list))
train_list = trainval_list[:train_count]
val_list = trainval_list[train_count:]
else:
train_count = int(train_ratio * total_file_count)
train_list = file_list[:train_count]
val_list = file_list[train_count:]
return train_list, val_list, test_list
def start(self, json_dir, save_dir, names=None, train_ratio=0.9):
images_dir, labels_dir = self.create_save_dir(save_dir)
if names is None or len(names) == 0:
print('class names will load from all json file')
names = self.get_class_names_from_all_json(json_dir)
print('find {} class names :'.format(len(names)), names)
if len(names) == 0:
return
self.save_list(names, os.path.join(save_dir, 'labels.txt'))
print('start convert')
all_json_list = []
for file in os.listdir(json_dir):
if not file.endswith('.json'):
continue
all_json_list.append(file)
train_list, val_list, test_list = self.split_train_val_test_dataset(all_json_list, train_ratio)
self.convert_dataset(json_dir, train_list, images_dir, labels_dir, names, 'train')
self.convert_dataset(json_dir, val_list, images_dir, labels_dir, names, 'val')
if __name__ == '__main__':
cm = ConvertManager()
cm.start(r'data_json', r'data_txt')
## 数据集划分 文件夹
import os, shutil, random
import numpy as np
TXT_path = 'data_txt' # 原TXT文件
Image_path = 'images' # 原图片文件
dataset_path = 'dataset/dataset1' # 保存的目标位置
val_size, test_size = 0.1, 0.2
os.makedirs(dataset_path, exist_ok=True)
os.makedirs(f'{dataset_path}/images', exist_ok=True)
os.makedirs(f'{dataset_path}/images/train', exist_ok=True)
os.makedirs(f'{dataset_path}/images/val', exist_ok=True)
os.makedirs(f'{dataset_path}/images/test', exist_ok=True)
os.makedirs(f'{dataset_path}/labels/train', exist_ok=True)
os.makedirs(f'{dataset_path}/labels/val', exist_ok=True)
os.makedirs(f'{dataset_path}/labels/test', exist_ok=True)
path_list = np.array([i.split('.')[0] for i in os.listdir(TXT_path) if 'txt' in i])
random.shuffle(path_list)
train_id = path_list[:int(len(path_list) * (1 - val_size - test_size))]
val_id = path_list[int(len(path_list) * (1 - val_size - test_size)):int(len(path_list) * (1 - test_size))]
test_id = path_list[int(len(path_list) * (1 - test_size)):]
for i in train_id:
shutil.copy(f'{Image_path}/{i}.jpg', f'{dataset_path}/images/train/{i}.jpg')
shutil.copy(f'{TXT_path}/{i}.txt', f'{dataset_path}/labels/train/{i}.txt')
for i in val_id:
shutil.copy(f'{Image_path}/{i}.jpg', f'{dataset_path}/images/val/{i}.jpg')
shutil.copy(f'{TXT_path}/{i}.txt', f'{dataset_path}/labels/val/{i}.txt')
for i in test_id:
shutil.copy(f'{Image_path}/{i}.jpg', f'{dataset_path}/images/test/{i}.jpg')
shutil.copy(f'{TXT_path}/{i}.txt', f'{dataset_path}/labels/test/{i}.txt')
还需要配置数据对应的yaml文件
path:这里看情况填，是根目录
train:train
val:val
test:test
names:
0:railway
1:person
三、模型训练
打开autodl租用实例，在vscode的SSHremote插件中输入账号密码登录到服务器，然后把数据和配置文件放到里面。我就用我本地的展示了。
requirements.txt文件中要修改一下
因为后面会报错，np.int等在numpy1.20后停用，np如果是1.20版则pandas版本也要改。
然后打开终端进入根目录seg安装环境。
pip install -r requirements.txt
这时候，检查yaml文件的根目录对不对就可以了
然后开始训练,dataset1.yaml是自己写的配置文件名称，放在data目录下，yolov5s-seg.pt放在根目录下，如果嫌默认300轮训练太久，可以加上–epoch 100，训练100轮
python segment/train.py --weights yolov5s-seg.pt --data data/dataset1.yaml
等待训练结束即可。
四、模型预测
把train-seg目录下的文件下载到自己电脑上，还有yaml文件也留存一份，拿出自己的原视频，放到data下images文件中，使用刚训练好的模型来进行预测。
python segment/predict.py --weights runs/train-seg/exp/weight/best.pt
检测完毕后，将train-seg下文件下载到自己电脑上。
展示下我的结果：
3157
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
