自定义数据集结构数据集主目录中包含以下几个文件夹分别是训练集的原图和mask 、测试集的原图和mask包含背景在内分割类别共4类: [bg, class1class2, class3]在mask中的像素值分别为064128192### 加载数据 def is_image_file(filename): #判断是否是图片 image_extensions [.jpg, .jpeg, .png, .gif, .bmp, .tiff, .tif, .webp, .svg, .ico] _, file_extension os.path.splitext(filename) file_extension file_extension.lower() return file_extension in image_extensions class Dataset(BaseDataset): # 数据集中用于图像分割的所有标签类别 CLASSES [bg, class1, class2, class3] def __init__( self, main_dirs, trainTrue, classesNone, augmentationNone, preprocessingNone, img_size1024, ): self.ids [] self.images_fps [] if train: for main_dir in main_dirs: images_dir os.path.join(main_dir, train_imgs) self.ids [file for file in os.listdir(images_dir) if is_image_file(file)] self.images_fps [os.path.join(images_dir, image_id) for image_id in [file for file in os.listdir(images_dir) if is_image_file(file)]] else: for main_dir in main_dirs: images_dir os.path.join(main_dir, val_imgs) self.ids [file for file in os.listdir(images_dir) if is_image_file(file)] self.images_fps [os.path.join(images_dir, image_id) for image_id in [file for file in os.listdir(images_dir) if is_image_file(file)]] # convert str names to class values on masks self.class_values [self.CLASSES.index(cls.lower()) for cls in classes] self.augmentation augmentation self.preprocessing preprocessing self.img_size img_size def __getitem__(self, i): # read data image_path self.images_fps[i] image cv2.imread(image_path) image cv2.cvtColor(image, cv2.COLOR_BGR2RGB) mask_path image_path.replace(_imgs, _segs) mask cv2.imread(mask_path, 0) mask // 64 # mask中的四个类别对应的值变为0 1 2 3 # 从标签中提取特定的类别 (可以传入需要的类别 我全选了 因为我全都要) masks [(mask v) for v in self.class_values] mask np.stack(masks, axis-1).astype(float) image cv2.resize(image, (self.img_size, self.img_size)) mask cv2.resize(mask, (self.img_size, self.img_size)) # # 图像增强应用 # if self.augmentation: # sample self.augmentation(imageimage, maskmask) # image, mask sample[image], sample[mask] # # 图像预处理应用 if self.preprocessing: sample self.preprocessing(imageimage, maskmask) image, mask sample[image], sample[mask] return image, mask def __len__(self): return len(self.ids) ### 图像增强 def get_training_augmentation(): train_transform [ albu.HorizontalFlip(p0.5), albu.ShiftScaleRotate(scale_limit0.5, rotate_limit0, shift_limit0.1, p1, border_mode0), albu.PadIfNeeded(min_height320, min_width320, always_applyTrue, border_mode0), albu.RandomCrop(height320, width320, always_applyTrue), albu.GaussNoise(p0.2), albu.Perspective(p0.5), albu.OneOf( [ albu.CLAHE(p1), albu.RandomBrightnessContrast(p1), albu.RandomGamma(p1), ], p0.9, ), albu.OneOf( [ albu.Sharpen(p1), albu.Blur(blur_limit3, p1), albu.MotionBlur(blur_limit3, p1), ], p0.9, ), albu.OneOf( [ albu.RandomBrightnessContrast(p1), albu.HueSaturationValue(p1), ], p0.9, ), ] return albu.Compose(train_transform) def get_validation_augmentation(): 调整图像使得图片的分辨率长宽能被32整除 test_transform [ albu.PadIfNeeded(384, 480) ] return albu.Compose(test_transform) def to_tensor(x, **kwargs): return x.transpose(2, 0, 1).astype(float32) def get_preprocessing(preprocessing_fn): 进行图像预处理操作 _transform [ albu.Lambda(imagepreprocessing_fn), albu.Lambda(imageto_tensor, maskto_tensor), ] return albu.Compose(_transform) # 创建模型并训练 if __name__ __main__: main_dirs [main_dir1, main_dir2, ...] epochs 200 img_size 1024 ENCODER resnet18 ENCODER_WEIGHTS imagenet CLASSES [bg, class1, class2, class3] ACTIVATION sigmoid DEVICE cuda # 使用DeepLabV3Plus模型 model smp.DeepLabV3Plus( encoder_nameENCODER, encoder_weightsENCODER_WEIGHTS, classeslen(CLASSES), activationACTIVATION, ) preprocessing_fn smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS) # 加载训练数据集 train_dataset Dataset( main_dirs main_dirs, trainTrue, augmentationget_training_augmentation(), preprocessingget_preprocessing(preprocessing_fn), classesCLASSES, img_size img_size, ) # 加载验证数据集 valid_dataset Dataset( main_dirs main_dirs, trainFalse, augmentationget_validation_augmentation(), preprocessingget_preprocessing(preprocessing_fn), classesCLASSES, img_sizeimg_size, ) # 需根据显卡的性能进行设置batch_size为每次迭代中一次训练的图片数num_workers为训练时的工作进程数如果显卡不太行或者显存空间不够将batch_size调低并将num_workers调为0 train_loader DataLoader(train_dataset, batch_size4, shuffleTrue, num_workers0) valid_loader DataLoader(valid_dataset, batch_size2, shuffleFalse, num_workers0) loss smp_utils.losses.DiceLoss() metrics [ smp_utils.metrics.IoU(threshold0.5), ] optimizer torch.optim.Adam([ dict(paramsmodel.parameters(), lr0.0001), ]) # 创建一个简单的循环用于迭代数据样本 train_epoch smp_utils.train.TrainEpoch( model, lossloss, metricsmetrics, optimizeroptimizer, deviceDEVICE, verboseTrue, ) valid_epoch smp_utils.train.ValidEpoch( model, lossloss, metricsmetrics, deviceDEVICE, verboseTrue, ) # 进行epochs轮次迭代的模型训练 max_score 0 for i in range(0, epochs): print(\nEpoch: {}.format(i)) train_logs train_epoch.run(train_loader) valid_logs valid_epoch.run(valid_loader) # 每次迭代保存下训练最好的模型 if max_score valid_logs[iou_score]: max_score valid_logs[iou_score] torch.save(model, ./best_model.pth) print(Model saved!) if i 25: optimizer.param_groups[0][lr] 1e-5 print(Decrease decoder learning rate to 1e-5!)训练的时候还遇到了问题显示报错ValueError: Expected more than 1 value per channel when training, got input size torch.Size([1, 256, 1, 1])当前批次仅包含一个样本batch size 1导致BatchNorm层无法正常工作因为BatchNorm需要每个通道有多于一个值来计算合理的均值和方差。segmentation_models_pytorch/utils/train.py文件中 修改代码如下如果要使用UNet、FPN、UnetPlusPlus等模型只要改成model smp.Unet or model smp.FPN or model smp.UnetPlusPlus()即可。