企业官网建设流程全解析

自定义数据集结构

数据集主目录中包含以下几个文件夹，分别是训练集的原图和mask 、测试集的原图和mask

包含背景在内，分割类别共4类: ["bg", "class1"，"class2", "class3"]

在mask中的像素值分别为0，64，128，192

### 加载数据 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, train=True, classes=None, augmentation=None, preprocessing=None, img_size=1024, ): 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(image=image, mask=mask) # image, mask = sample['image'], sample['mask'] # # 图像预处理应用 if self.preprocessing: sample = self.preprocessing(image=image, mask=mask) image, mask = sample['image'], sample['mask'] return image, mask def __len__(self): return len(self.ids) ### 图像增强 def get_training_augmentation(): train_transform = [ albu.HorizontalFlip(p=0.5), albu.ShiftScaleRotate(scale_limit=0.5, rotate_limit=0, shift_limit=0.1, p=1, border_mode=0), albu.PadIfNeeded(min_height=320, min_width=320, always_apply=True, border_mode=0), albu.RandomCrop(height=320, width=320, always_apply=True), albu.GaussNoise(p=0.2), albu.Perspective(p=0.5), albu.OneOf( [ albu.CLAHE(p=1), albu.RandomBrightnessContrast(p=1), albu.RandomGamma(p=1), ], p=0.9, ), albu.OneOf( [ albu.Sharpen(p=1), albu.Blur(blur_limit=3, p=1), albu.MotionBlur(blur_limit=3, p=1), ], p=0.9, ), albu.OneOf( [ albu.RandomBrightnessContrast(p=1), albu.HueSaturationValue(p=1), ], p=0.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(image=preprocessing_fn), albu.Lambda(image=to_tensor, mask=to_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_name=ENCODER, encoder_weights=ENCODER_WEIGHTS, classes=len(CLASSES), activation=ACTIVATION, ) preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS) # 加载训练数据集 train_dataset = Dataset( main_dirs = main_dirs, train=True, augmentation=get_training_augmentation(), preprocessing=get_preprocessing(preprocessing_fn), classes=CLASSES, img_size = img_size, ) # 加载验证数据集 valid_dataset = Dataset( main_dirs = main_dirs, train=False, augmentation=get_validation_augmentation(), preprocessing=get_preprocessing(preprocessing_fn), classes=CLASSES, img_size=img_size, ) # 需根据显卡的性能进行设置，batch_size为每次迭代中一次训练的图片数，num_workers为训练时的工作进程数，如果显卡不太行或者显存空间不够，将batch_size调低并将num_workers调为0 train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=0) valid_loader = DataLoader(valid_dataset, batch_size=2, shuffle=False, num_workers=0) loss = smp_utils.losses.DiceLoss() metrics = [ smp_utils.metrics.IoU(threshold=0.5), ] optimizer = torch.optim.Adam([ dict(params=model.parameters(), lr=0.0001), ]) # 创建一个简单的循环，用于迭代数据样本 train_epoch = smp_utils.train.TrainEpoch( model, loss=loss, metrics=metrics, optimizer=optimizer, device=DEVICE, verbose=True, ) valid_epoch = smp_utils.train.ValidEpoch( model, loss=loss, metrics=metrics, device=DEVICE, verbose=True, ) # 进行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()即可。