当前位置：首页 > news >正文

u2net网络模型训练自己数据集

news 2025/7/19 0:26:32

单分类
下载项目源码
项目源码

准备数据集
将json转为mask
json_to_dataset.py

import cv2
import json
import numpy as np
import os
import sys
import globdef func(file):with open(file, mode='r', encoding="utf-8") as f:configs = json.load(f)shapes = configs["shapes"]png_class = np.zeros((configs["imageHeight"], configs["imageWidth"], 1), np.uint8)png_other = np.zeros((configs["imageHeight"], configs["imageWidth"], 1), np.uint8)for shape in shapes:label = shape['label']if label == 'class':cv2.fillPoly(png_class, [np.array(shape["points"], np.int32)], (255))else:cv2.fillPoly(png_other, [np.array(shape["points"], np.int32)], (255))png = png_class - png_otherreturn pngif __name__ == "__main__":json_dir = "image"save_dir = 'image/masks'for file in os.listdir(json_dir):print('***************', file)if file.endswith(".json"):# 在这里添加后续步骤png = func(os.path.join(json_dir, file))print(png.shape)save_path = save_dir + '/' + os.path.splitext(file)[0] + ".png"cv2.imwrite(save_path, png)print('***************', save_path)

创建文件路径

train_data|__DUTS|__DUTS-TR|__im_aug			# mask图|__gt_aug			# 原图|__DUTS-TE|__im_aug			# mask图|__gt_aug			# 原图

训练
修改u2net_train.py

# ------- 2. set the directory of training dataset --------model_name = 'u2net' #'u2netp'											# 选择模型# 修改为对应的数据集路径
data_dir = os.path.join(os.getcwd(), 'train_data' + os.sep)
tra_image_dir = os.path.join('DUTS', 'DUTS-TR', 'im_aug' + os.sep)
tra_label_dir = os.path.join('DUTS', 'DUTS-TR', 'gt_aug' + os.sep)image_ext = '.jpg'
label_ext = '.png'model_dir = os.path.join(os.getcwd(), 'saved_models', model_name + os.sep)

alpha_export.py为转换模型，u2net_test.py模型预测

多分类
准备数据集
将json转为mask，由于是多分类，每一个分类一个rgb值，例如分类1（1，1，1），分类2（2，2，2），分类3（3，3，3）；

import cv2
import json
import numpy as np
import os
import sys
import globdef func(file):with open(file, mode='r', encoding="utf-8") as f:configs = json.load(f)shapes = configs["shapes"]class1 = np.zeros((configs["imageHeight"], configs["imageWidth"], 3), np.uint8)class2 = np.zeros((configs["imageHeight"], configs["imageWidth"], 3), np.uint8)class3 = np.zeros((configs["imageHeight"], configs["imageWidth"], 3), np.uint8)for shape in shapes:label = shape['label']if label == 'class1':cv2.fillPoly(class1, [np.array(shape["points"], np.int32)], (1, 1, 1))elif label == 'class2':cv2.fillPoly(class2, [np.array(shape["points"], np.int32)], (2, 2, 2))elif label == 'class3':cv2.fillPoly(class3, [np.array(shape["points"], np.int32)], (3, 3, 3))return class1 + class2 + class3if __name__ == "__main__":json_dir = "./train_data/labels_json"save_dir = './train_data/masks'for file in os.listdir(json_dir):print('***************', file)if file.endswith(".json"):# 在这里添加后续步骤png = func(os.path.join(json_dir, file))print(png.shape)save_path = save_dir + '/' + os.path.splitext(file)[0] + ".png"cv2.imwrite(save_path, png)print('***************', save_path)

数据集文件夹

datasets|__train_data|__images		# 原图|__1.jpg|__2.jpg|__3.jpg|__masks		# mask图|__1.png|__2.png|__3.png|__test_data|__images|__1.jpg|__2.jpg|__3.jpg|__masks|__1.png|__2.png|__3.png|__my_results_2|__predeict_lables|__predict_masks

训练

# ------- 2. set the directory of training process --------
model_name = 'u2net'  # 'u2net'							# 选择模型
model_dir = os.path.join('saved_models', model_name + os.sep)# 图片的文件类型
image_ext = '.jpg'
label_ext = '.png'# train阶段
# 数据集路径
data_dir = os.path.join("datasets/train_data" + os.sep)
# 原始图片路径
tra_image_dir = os.path.join('images' + os.sep)
# 图片的标签路径
tra_label_dir = os.path.join('masks' + os.sep)# val阶段
# 数据集类别
num_classes = 4			# 背景+类别数
name_classes = ["background", "class1", "class2","class3"]
# 原始图片路径
images_path = "datasets/test_data/images/"
# 图片的标签路径
gt_dir = "datasets/test_data/masks/"
# 存放推理结果图片的路径
pred_dir = "datasets/test_data/predict_masks/"
predict_label = "datasets/test_data/predict_labels/"
# 存放 miou 计算结果的 图片
miou_out_path = "miou_out_tab"
# 预训练权重
seg_pretrain_u2netp_path = 'saved_models/pretrain_model/segm_u2net.pth'
if not os.path.exists(model_dir):os.makedirs(model_dir)epoch_num = 100
batch_size = 2
train_num = 0
val_num = 0

执行u2net_train.py，开始训练

预测

if __name__ == "__main__":#   miou_mode用于指定该文件运行时计算的内容#   miou_mode为0代表整个miou计算流程，包括获得预测结果、计算miou。#   miou_mode为1代表仅仅获得预测结果。#   miou_mode为2代表仅仅计算miou。#   分类个数+1、如2+1# num_classes = 3# name_classes = ["background", "green", "red"]num_classes = 4name_classes =["background", "class1", "class2","class3"]# 原始图片路径images_path = "datasets_ButtonCell/test_data/images/"# 图片的标签路径gt_dir = "datasets_ButtonCell/test_data/masks/"# 存放推理结果图片的路径pred_dir = "datasets_ButtonCell/test_data/predict_masks/"predict_label = "datasets_ButtonCell/test_data/predict_labels/"# 存放 miou 计算结果的 图片miou_out_path = "miou_out"# 模型路径model_dir = './saved_models/u2netp/u2netp.pth'eval_print_miou(num_classes, name_classes, images_path, gt_dir, pred_dir, predict_label, miou_out_path, model_dir)