语义分割任务中VOC数据集的制作,任务中只有一种标签:gas
文章目录
- 1、由黑白图像识别为txt标签
- 2、txt转json
- 3、数据集转VOC格式
1、由黑白图像识别为txt标签
由于使用CycleGAN网络进行风格迁移学习,生成了大量伪标签图像,因此需要自动提取出标签,标签格式为YOLOv8训练所需要的语义分割的标签格式
例如:提取出下图中白色的区域的标签信息
txt标签:
import os
import cv2
import numpy as np
'''
该段代码读取文件夹中的二值图像,输出白色区域的标注信息
格式为YOLOv8图像分割的txt格式
'''
def process_images(input_folder, output_folder):# 如果输出文件夹不存在,则创建它if not os.path.exists(output_folder):os.makedirs(output_folder)# 列出输入文件夹中的所有图像文件image_files = [f for f in os.listdir(input_folder) if f.endswith(('.png', '.jpg', '.jpeg', '.bmp'))]total_images = len(image_files)print(f"在文件夹'{input_folder}'中找到了 {total_images} 张图像。")# 处理每张图像for idx, filename in enumerate(image_files, 1):print(f"正在处理图像 {idx} / {total_images} : {filename}")# 加载图像image_path = os.path.join(input_folder, filename)image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)# 检查图像是否正确加载if image is None:print(f"图像 {filename} 未能正确加载。")continue# 对图像进行阈值处理以获得二值图像_, binary_image = cv2.threshold(image, 128, 255, cv2.THRESH_BINARY)# 寻找轮廓contours, _ = cv2.findContours(binary_image, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)# 初始化一个列表来存储轮廓坐标contour_coordinates = []# 获取图像的尺寸以进行归一化height, width = image.shape# 遍历轮廓以提取归一化坐标for contour in contours:# 初始化一个列表来存储单个轮廓的坐标coords = ['0'] # 开头的数字0(没有小数点)for point in contour:# 归一化x和y坐标并添加到列表中x_normalized = point[0][0] / widthy_normalized = point[0][1] / heightcoords.extend([f"{x_normalized:.6f}", f"{y_normalized:.6f}"])# 将单个轮廓的坐标添加到主列表中contour_coordinates.append(' '.join(coords))# 定义输出文本文件的路径output_file_path = os.path.join(output_folder, os.path.splitext(filename)[0] + '.txt')# 将坐标写入文本文件with open(output_file_path, 'w') as file:file.write('\n'.join(contour_coordinates))
# 输入和输出文件夹路径(请替换为实际路径)
input_folder_path = 'D:\Desktop\\test0326' # 替换为实际的输入图片文件夹
output_folder_path = 'D:\Desktop\\test0326' # 替换为实际的输出txt标签文件夹
# 调用函数以处理图像
process_images(input_folder_path, output_folder_path)
2、txt转json
由txt文件转换为labelme标注的json格式,本人的labelme版本为5.3.1,json文件内涵图像数据,具体格式为:
json_data = {"version": "5.3.1","flags": {},"shapes": [],"imagePath": img_filename,"imageData": encoded_string,"imageHeight": height,"imageWidth": width}
代码如下:
import os
import json
import base64
from PIL import Image
import io# 文件夹路径
label_folder = 'D:\\Desktop\\Unet\\dataset\\500txt' # txt标签文件夹
image_folder = 'D:\\Desktop\\Unet\\dataset\\500images' # 图像文件夹
output_folder = 'D:\\Desktop\\Unet\\dataset\\500jsons' # json标签输出的文件夹# 确保输出文件夹存在
os.makedirs(output_folder, exist_ok=True)
# 遍历文件夹中的所有 txt 文件
for filename in os.listdir(label_folder):if filename.endswith('.txt'):# 构造完整的文件路径txt_path = os.path.join(label_folder, filename)# 构造对应的图像文件路径img_filename = filename.replace('.txt', '.jpg')img_path = os.path.join(image_folder, img_filename)# 读取图像尺寸并转换为 base64with Image.open(img_path) as img:width, height = img.sizebuffered = io.BytesIO()img.save(buffered, format="JPEG")encoded_string = base64.b64encode(buffered.getvalue()).decode('utf-8')# 初始化 JSON 数据结构json_data = {"version": "5.3.1","flags": {},"shapes": [],"imagePath": img_filename,"imageData": encoded_string,"imageHeight": height,"imageWidth": width}# 读取 txt 文件并解析多边形坐标with open(txt_path, 'r') as file:for line in file:parts = line.strip().split()label_index = parts[0] # 假设第一部分是类别索引if len(parts) < 3 or len(parts) % 2 == 0:print(f"Unexpected format in {filename}")continue# 将相对坐标转换为实际坐标(像素值)label = 'gas' if label_index == '0' else 'unknown'points = [(float(parts[i]) * width, float(parts[i + 1]) * height) for i in range(1, len(parts), 2)]shape_data = {"label": label,"points": points,"group_id": None,"description": "","shape_type": "polygon","flags": {}}json_data["shapes"].append(shape_data)# 输出 JSON 文件output_path = os.path.join(output_folder, filename.replace('.txt', '.json'))with open(output_path, 'w') as jsonfile:json.dump(json_data, jsonfile, indent=2)print(f"文件 {output_path} 已转换完成!")print("Conversion to JSON completed.")
3、数据集转VOC格式
准备好原图和json格式的标签后,转换为VOC所需的语义分割数据集
还需要准备一个labels.txt文件,里面写上所有的类别,格式如下:
__ignore__
_background_
gas
由于我的任务中只有gas这一类,所以只需要写gas,如果是多类,则依次往下添加即可,“ignore”和“background”这两行不要动,准备好labels.txt文件后还需要将路径填到下方转换代码中注明的位置:
from __future__ import print_function
import argparse
import glob
import os
import os.path as osp
import sys
import imgviz
import numpy as np
import labelmedef main(args):if osp.exists(args.output_dir):print("Output directory already exists:", args.output_dir)sys.exit(1)os.makedirs(args.output_dir)os.makedirs(osp.join(args.output_dir, "JPEGImages"))os.makedirs(osp.join(args.output_dir, "SegmentationClassnpy"))os.makedirs(osp.join(args.output_dir, "SegmentationClass"))if not args.noviz:os.makedirs(osp.join(args.output_dir, "SegmentationClassVisualization"))print("Creating dataset:", args.output_dir)class_names = []class_name_to_id = {}for i, line in enumerate(open(args.labels).readlines()):class_id = i - 1 # starts with -1class_name = line.strip()class_name_to_id[class_name] = class_idif class_id == -1:assert class_name == "__ignore__"continueelif class_id == 0:assert class_name == "_background_"class_names.append(class_name)class_names = tuple(class_names)print("class_names:", class_names)out_class_names_file = osp.join(args.output_dir, "class_names.txt")with open(out_class_names_file, "w") as f:f.writelines("\n".join(class_names))print("Saved class_names:", out_class_names_file)for filename in glob.glob(osp.join(args.input_dir, "*.json")):print("Generating dataset from:", filename)label_file = labelme.LabelFile(filename=filename)base = osp.splitext(osp.basename(filename))[0]out_img_file = osp.join(args.output_dir, "JPEGImages", base + ".jpg")out_lbl_file = osp.join(args.output_dir, "SegmentationClassnpy", base + ".npy")out_png_file = osp.join(args.output_dir, "SegmentationClass", base + ".png")if not args.noviz:out_viz_file = osp.join(args.output_dir,"SegmentationClassVisualization",base + ".jpg",)with open(out_img_file, "wb") as f:f.write(label_file.imageData)img = labelme.utils.img_data_to_arr(label_file.imageData)lbl, _ = labelme.utils.shapes_to_label(img_shape=img.shape,shapes=label_file.shapes,label_name_to_value=class_name_to_id,)labelme.utils.lblsave(out_png_file, lbl)np.save(out_lbl_file, lbl)if not args.noviz:viz = imgviz.label2rgb(label=lbl,# img=imgviz.rgb2gray(img),# img=img,font_size=15,label_names=class_names,loc="rb",)imgviz.io.imsave(out_viz_file, viz)def get_args():parser = argparse.ArgumentParser()parser.add_argument("--input_dir", default="D:\Desktop\\Unet\dataset\\500jsons", type=str, ## 在这里输入json标签文件夹help="input annotated directory")parser.add_argument("--output_dir", default="D:\Desktop\\Unet\dataset\VOC500images\\", type=str, ## 在这里输入图片文件夹 help="output dataset directory")parser.add_argument("--labels", default="D:\Desktop\\Unet\\dataset\\labels.txt", type=str, ## 输入labels.txt文件所在的路径help="labels file")parser.add_argument("--noviz", help="no visualization", action="store_true")args = parser.parse_args()return args
if __name__ == "__main__":args = get_args()main(args)