本文主要是介绍学习系列三:V8目标检测与分割自动化标注,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!
学习系列三:YOLOv8目标检测与分割自动化标注
提示:本次文章主要介绍yolov8目标检测与自动化标注(较简单,通用性比较强,标签格式txt),yolov8实例分割与自动化标注(程序较复杂,自动化标注效果有待提升,标签格式txt),大家有更好的想法可以在评论区进行交流。
文章目录
- 学习系列三:YOLOv8目标检测与分割自动化标注
- 一、YOLOv8训练目标检测数据集
- 1.1 数据集的划分
- 1.2 目标检测训练脚本
- 二、实例化推理检测函数
- 三、目标检查自动化标注与标签生成(txt格式)
- 四、YOLOv8训练实例分割数据集
- 五、实例化推理检测函数
- 六、自动化标注与标签生成
一、YOLOv8训练目标检测数据集
1.1 数据集的划分
在进行训练目标检测数据前,需要划分训练集、验证集和测试集,可以进行人工的划分,可以使用代码进行自主进行划分,这里主要讲述使用代码进行自主划分训练集、验证集和测试集。
在划分之前需要看一下数据集的格式:数据集文件夹(tuoluo-aug), 数据集图片文件夹(images),数据集标签文件夹(labels, txt格式标签文件)
代码进行自主划分训练集、验证集和测试集。这里可以自己定义训练集、验证集和测试集所占比例多少。
from ultralytics.data.utils import autosplitautosplit(path='/home/xiao/dataset/tuoluo-aug/images',weights=(0.8, 0.2, 0.00), # (train, validation, test) fractional splitsannotated_only=False # split only images with annotation file when True
)
运行之后在 tuoluo-aug文件下会生成txt文件,如下:
1.2 目标检测训练脚本
针对YOLOv8这里我写一个简单的脚本训练目标检测数据集,放在工程ultralytics-xiao下面,脚本文件名称定义为tuoluo.py,代码如下:
from ultralytics import YOLOmodel = YOLO("/home/xiao/ultralytics-main/weights/yolov8n.pt")
model.train(data="/home/xiao/ultralytics-main/ultralytics/cfg/datasets/bolt-detachment-aug.yaml",epochs=120,imgsz = 1280,device= [0],workers = 2,batch =8,patience=120)metrics = model.val() # 在验证集上评估模型性能
在上述代码中数据集配置文件路径:“/home/xiao/ultralytics-main/ultralytics/cfg/datasets/bolt-detachment-aug.yaml”
配置文件内容为:
path: /home/xiao/dataset/tuoluo-aug/
train: autosplit_train.txt
val: autosplit_val.txt
test: autosplit_test.txt names:0: detachmentnc : 1
上述完成后,运行tuoluo.py即可开启训练。(上述数据集标签为:detachment)
二、实例化推理检测函数
针对上述训练完成后,会生成对应最好的权重,试用该权重进行推理新的图片,代码如下(封装成类函数了):
from ultralytics import YOLO
import cv2
import numpy as np
import time
import os# 加载模型
model_file = './weight/best.pt'model = YOLO(model_file)
objs_labels = model.names
print(objs_labels)class yolo_demo:def __init__(self):self.model_file_1 = './weight/best.pt'self.model_1 = YOLO(self.model_file_1)self.objs_labels_1 = self.model_1.namesdef img_file_folder(self, img_folder):image_paths = []for image_name in os.listdir(img_folder):image_path = os.path.join(img_folder, image_name)if os.path.isfile(image_path) and image_name.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff')):image_paths.append(image_path)return image_pathsdef main(self, img_folder):for i, img_path in enumerate(self.img_file_folder(img_folder)):# 获取图片文件名img_id = img_path.split(os.sep)[-1].split('.')[0]# bgrimage_bgr = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), -1)result = list(self.model_1(image_bgr, stream=True, conf=0.7))[0]boxes = result.boxesboxes = boxes.cpu().numpy()for box in boxes.data:l, t, r, b = box[:4].astype(np.int32)conf, id = box[4:]if id == 0:# 绘制标签矩形框cv2.rectangle(image_bgr, (l, t), (r, b), (0, 255, 0), 1)label = "detachment"cv2.putText(image_bgr, label, (l, t - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)output_image_path = f"./output-v4/{img_id}.jpg"cv2.imwrite(output_image_path, image_bgr)if __name__ =="__main__":plate_demo = yolo_demo()img_path = r"./detachment-images"
三、目标检查自动化标注与标签生成(txt格式)
适用于已经使用少量数据集图片进行YOLOv8训练,得到权重并使用该权重对新的数据集图片进行自动标注,后期可人工进行检查标注的好坏,进行简单修改,相比较于从头进行标注节省了大量的时间。
"""
yolov8 自动标注, 后续需要手工进行修改校准
"""
from ultralytics import YOLO
import cv2
import os
import numpy as np
import matplotlib.pyplot as plt
import glob
import shutil
import tqdm# 加载模型
model_file = "./weights1to4/detachment.pt"
model = YOLO(model_file)
objs_labels = model.names
print(objs_labels)# 读取图片
images_list = glob.glob('./img-detachment-9/*.jpg')print(len(images_list))# 创建标签文件夹
if not os.path.exists('./detachment-labels'):os.mkdir('./detachment-labels')# 标注
def image_2_yolo():# 保存路径前缀savePathPrefix = "./detachment-labels/"# 遍历每张图片for img in images_list:# 获取图片文件名img_id = img.split(os.sep)[-1].split('.')[0]# 读取图片img_data = cv2.imread(img)# 检测result = list(model(img_data, stream=True, conf=0.5))[0]boxes = result.boxesboxes = boxes.cpu().numpy()yolo_boxes = []# 获取图片宽高img_h, img_w, _ = img_data.shape# 遍历每个框for box in boxes.data:l, t, r, b = box[:4].astype(np.int32)conf, id = box[4:]# 筛选出detachment类别, 转为yolo格式:类别id, x_center, y_center, width, height, 归一化到0-1, 保留6位小数if id == 0:class_label = int(id)x_center = round((l + r) / 2 / img_w, 6)y_center = round((t + b) / 2 / img_h, 6)width = round((r - l) / img_w, 6)height = round((b - t) / img_h, 6)yolo_boxes.append([class_label, x_center, y_center, width, height])# 写入txt文件# 生成yolo格式的标注文件yoloLabelFile = savePathPrefix + img_id + '.txt'with open(yoloLabelFile, 'w') as f:for yolo_box in yolo_boxes:f.write(' '.join([str(i) for i in yolo_box]) + '\n')if __name__ == '__main__':image_2_yolo()
四、YOLOv8训练实例分割数据集
from ultralytics import YOLOmodel = YOLO('/home/hxz/xiao/ultralytics-xiao/ultralytics/cfg/models/v8/yolov8-seg.yaml').load('/home/hxz/xiao/ultralytics-xiao/runs/segment/train6/weights/best.pt')model.train(data="/home/hxz/xiao/ultralytics-xiao/ultralytics/cfg/datasets/seg-bolt-line-cz933-weitiao.yaml",task="segment",mode="train",overlap_mask=False,batch=16,device=0,epochs=600,patience=600,imgsz=128)metrics = model.val(iou=0.7)try:a = metrics.box.map # map50-95b = metrics.box.map50 # map50c = metrics.box.map75 # map75d = metrics.box.maps # 返回一个列表,包含每个类别的mAP值(IoU在0.5到0.95)print("result:", a, b, c, d)except Exception as e:print('f{有一些问题}')
五、实例化推理检测函数
from ultralytics import YOLO# Load a model
model = YOLO("/home/hxz/xiao/ultralytics-xiao/runs/segment/train6/weights/best.pt") # load a custom model# Predict with the model
source = "/home/hxz/xiao/ultralytics-xiao/test-img" # predict on an imageresults = model.predict(source, save=True)
六、自动化标注与标签生成
这部分是根据比较少的数据集训练得到的数据的权重,根据权重区预测新的图片,获得mask,根据mask读取坐标保存至txt文件。
from ultralytics import YOLO
import numpy as np
import cv2
import os
import matplotlib.pyplot as plt
import glob# 加载分割模型
model = YOLO("./weights/fks-seg.pt") # 加载自定义模型# 预测源
image_folder = r"./fks-seg/images"# 获取文件夹中所有图片的路径
image_paths = glob.glob(os.path.join(image_folder, "*.jpg"))# 保存分割结果坐标点的目录
output_dir = "./fks-seg/txt_labels"
os.makedirs(output_dir, exist_ok=True)for image_path in image_paths:results = model.predict(image_path)base_name = os.path.basename(image_path).split('.')[0]# 创建输出文件路径txt_filename = os.path.join(output_dir, f"{base_name}.txt")for result in results:masks = result.masks # 获取分割掩码classes = result.boxes.cls if result.boxes is not None else [] # 获取每个掩码的类别# 获取图像的宽度和高度以进行归一化height, width = result.orig_img.shape[:2]print(height, width)# 创建输出文件路径txt_filename = os.path.join(output_dir, f"{os.path.basename(result.path).split('.')[0]}.txt")with open(txt_filename, 'w') as f:if masks is not None:for i, mask in enumerate(masks.data):# 将 PyTorch 张量转换为 NumPy 数组并确保掩码是二值图像mask = mask.cpu().numpy().astype(np.uint8)contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)# 获取当前掩码对应的类别class_id = int(classes[i]) if i < len(classes) else -1# 准备一个列表来存储该类别的所有归一化坐标coordinates = []for contour in contours:# 将轮廓的坐标点进行归一化normalized_coordinates = []for point in contour:x, y = point[0]x_normalized = x / mask.shape[1]y_normalized = y / mask.shape[0]normalized_coordinates.append(f"{x_normalized:.6f} {y_normalized:.6f}")# 获取当前掩码对应的类别class_id = int(classes[i]) if i < len(classes) else -1# 写入 YOLO 格式的标签文件(类别 + 轮廓坐标)if normalized_coordinates:f.write(f"{class_id} " + " ".join(normalized_coordinates) + "\n")print(f"Saved YOLO segmentation coordinates to {txt_filename}")
将生成的txt文件转json后,把图片和生成的json文件放在同一个文件夹,可在lableme中查看自动标注的情况。
import os
import json
import cv2
import numpy as np
import globdef yolo_to_labelme(txt_file, img_file, class_names, output_json):if not os.path.exists(txt_file):print(f"File {txt_file} does not exist.")returnimg = cv2.imdecode(np.fromfile(img_file, dtype=np.uint8), -1)height, width = img.shape[:2]shapes = []with open(txt_file, 'r') as f:lines = f.readlines()for line in lines:parts = line.strip().split()class_id = int(parts[0])# 将归一化的坐标转换为非归一化的像素坐标points = [(float(parts[i]) * width, float(parts[i + 1]) * height) for i in range(1, len(parts), 2)]shape = {"label": class_names[class_id],"points": points,"group_id": None,"shape_type": "polygon","flags": {}}shapes.append(shape)labelme_data = {"version": "4.5.6","flags": {},"shapes": shapes,"imagePath": os.path.basename(img_file),"imageData": None,"imageHeight": height,"imageWidth": width}with open(output_json, 'w') as f:json.dump(labelme_data, f, indent=4)print(f"Saved JSON to {output_json}")def batch_process_yolo_to_labelme(image_folder, label_folder, output_folder, class_names):os.makedirs(output_folder, exist_ok=True)# 获取所有的图片文件image_files = glob.glob(os.path.join(image_folder, "*.jpg"))for img_file in image_files:# 获取对应的 txt 文件base_name = os.path.basename(img_file).split('.')[0]txt_file = os.path.join(label_folder, f"{base_name}.txt")# 输出的 JSON 文件路径output_json = os.path.join(output_folder, f"{base_name}.json")# 转换 YOLO 到 LabelMe 格式yolo_to_labelme(txt_file, img_file, class_names, output_json)# 示例使用
class_names = ["red"] # 根据你的类别顺序填写
image_folder = r"./fks-seg/images"
label_folder = r"./fks-seg/txt_labels"
output_folder = r"./fks-seg/json_labels"batch_process_yolo_to_labelme(image_folder, label_folder, output_folder, class_names)
这篇关于学习系列三:V8目标检测与分割自动化标注的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!