UNet训练道路裂缝检测数据集6572多张【裂缝语义分割数据集】unet unet fcn segformer deeplab等语义分割算法道路裂缝检测数据集数据集概述数据来源该数据集由10个子数据集预处理和大小调整为400x400具体子数据集包括Crack500 Deepcrack Sdnet Cracktree Gaps Volker Rissbilder Noncrack Massimo Ceramic数据规模共6000多张图像每张图像都有对应的像素级标签。图像格式JPG标签格式二值标签有两种格式0-255像素0-1像素用途适用于UNet、UNet、FCN、SegFormer、DeepLab等语义分割算法。数据集结构深色版本road_crack_detection/ ├── dataset/ │ ├── images/ │ │ └── *.jpg │ ├── labels_0_255/ │ │ └── *.png │ ├── labels_0_1/ │ │ └── *.png ├── models/ │ └── unet/ │ └── unetpp/ │ └── fcn/ │ └── segformer/ │ └── deeplab/ ├── src/ │ ├── train.py │ ├── predict.py │ ├── utils.py │ ├── dataset.py ├── weights/ │ └── best_model.pth ├── requirements.txt └── README.md安装依赖首先确保你已经安装了必要的库。创建一个requirements.txt文件内容如下深色版本torchtorchvisionnumpypandasmatplotlibtqdmpyyamlopencv-python然后使用以下命令安装依赖bash深色版本pip install -r requirements.txt2. 数据集准备确保你的数据集已经按照以下结构组织深色版本dataset/├── images/│ └── *.jpg├── labels_0_255/│ └── *.png├── labels_0_1/│ └── *.png每个文件夹中包含对应的图像文件和标签文件。确保所有图像文件都是.jpg格式标签文件是PNG格式。数据集类创建一个数据集类用于加载和预处理数据。3.1 src/dataset.pypython深色版本import osimport torchfrom torch.utils.data import Datasetfrom torchvision import transformsfrom PIL import Imageclass RoadCrackDataset(Dataset):definit(self, image_dir, label_dir, transformNone, label_type‘0_1’):self.image_dir image_dirself.label_dir label_dirself.transform transformself.image_files os.listdir(image_dir)self.label_type label_typedef __len__(self): return len(self.image_files) def __getitem__(self, index): img_path os.path.join(self.image_dir, self.image_files[index]) label_path os.path.join(self.label_dir, self.image_files[index].replace(.jpg, .png)) image Image.open(img_path).convert(RGB) label Image.open(label_path).convert(L) if self.label_type 0_255: label np.array(label) // 255 else: label np.array(label) if self.transform: image self.transform(image) label torch.from_numpy(label).long() return image, labeldef get_data_loaders(image_dir, label_dir, batch_size16, num_workers4, label_type‘0_1’):transform transforms.Compose([transforms.Resize((400, 400)),transforms.ToTensor(),transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]),])dataset RoadCrackDataset(image_dir, label_dir, transformtransform, label_typelabel_type) train_size int(0.8 * len(dataset)) val_size len(dataset) - train_size train_dataset, val_dataset torch.utils.data.random_split(dataset, [train_size, val_size]) train_loader torch.utils.data.DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue, num_workersnum_workers) val_loader torch.utils.data.DataLoader(val_dataset, batch_sizebatch_size, shuffleFalse, num_workersnum_workers) return train_loader, val_loader模型定义这里以UNet为例定义模型并进行训练。4.1 src/train.pypython深色版本import torchimport torch.optim as optimfrom torch.utils.tensorboard import SummaryWriterfrom tqdm import tqdmfrom src.dataset import get_data_loadersimport torch.nn as nnimport torch.nn.functional as Ffrom models.unet import UNetdef train_model(image_dir, label_dir, epochs100, batch_size16, learning_rate1e-4, label_type‘0_1’):device torch.device(“cuda” if torch.cuda.is_available() else “cpu”)model UNet(n_channels3, n_classes2) model model.to(device) train_loader, val_loader get_data_loaders(image_dir, label_dir, batch_sizebatch_size, label_typelabel_type) optimizer optim.Adam(model.parameters(), lrlearning_rate) criterion nn.CrossEntropyLoss() writer SummaryWriter() for epoch in range(epochs): model.train() running_loss 0.0 for images, labels in tqdm(train_loader, descfEpoch {epoch 1}/{epochs}): images images.to(device) labels labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() train_loss running_loss / len(train_loader) writer.add_scalar(Training Loss, train_loss, epoch) model.eval() running_val_loss 0.0 with torch.no_grad(): for images, labels in val_loader: images images.to(device) labels labels.to(device) outputs model(images) loss criterion(outputs, labels) running_val_loss loss.item() val_loss running_val_loss / len(val_loader) writer.add_scalar(Validation Loss, val_loss, epoch) print(fEpoch {epoch 1}/{epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}) torch.save(model.state_dict(), weights/best_model.pth) writer.close()ifname “main”:image_dir “dataset/images”label_dir “dataset/labels_0_1”train_model(image_dir, label_dir)5. 模型评估训练完成后可以通过验证集和测试集来评估模型的性能。示例如下5.1 src/predict.pypython深色版本import torchimport matplotlib.pyplot as pltfrom models.unet import UNetfrom src.dataset import get_data_loadersimport numpy as npdef predict_and_plot(image_dir, label_dir, model_path, num_samples5, label_type‘0_1’):device torch.device(“cuda” if torch.cuda.is_available() else “cpu”)model UNet(n_channels3, n_classes2) model.load_state_dict(torch.load(model_path)) model model.to(device) model.eval() _, val_loader get_data_loaders(image_dir, label_dir, label_typelabel_type) fig, axes plt.subplots(num_samples, 3, figsize(15, 5 * num_samples)) with torch.no_grad(): for i, (images, labels) in enumerate(val_loader): if i num_samples: break images images.to(device) labels labels.to(device) outputs model(images) predictions torch.argmax(outputs, dim1).cpu().numpy() labels labels.cpu().numpy() images images.cpu().numpy().transpose((0, 2, 3, 1)) for j in range(len(images)): ax axes[j] if num_samples 1 else axes ax[0].imshow(images[j]) ax[0].set_title(Input Image) ax[0].axis(off) ax[1].imshow(labels[j], cmapgray) ax[1].set_title(True Label) ax[1].axis(off) ax[2].imshow(predictions[j], cmapgray) ax[2].set_title(Predicted Label) ax[2].axis(off) plt.tight_layout() plt.show()ifname “main”:image_dir “dataset/images”label_dir “dataset/labels_0_1”model_path “weights/best_model.pth”predict_and_plot(image_dir, label_dir, model_path)6. 模型定义UNet6.1 models/unet.pypython深色版本import torchimport torch.nn as nnimport torch.nn.functional as Fclass DoubleConv(nn.Module):definit(self, in_channels, out_channels):super(DoubleConv, self).init()self.conv nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size3, padding1),nn.BatchNorm2d(out_channels),nn.ReLU(inplaceTrue),nn.Conv2d(out_channels, out_channels, kernel_size3, padding1),nn.BatchNorm2d(out_channels),nn.ReLU(inplaceTrue))def forward(self, x): return self.conv(x)class UNet(nn.Module):definit(self, n_channels, n_classes):super(UNet, self).init()self.n_channels n_channelsself.n_classes n_classesself.inc DoubleConv(n_channels, 64) self.down1 nn.Sequential( nn.MaxPool2d(2), DoubleConv(64, 128) ) self.down2 nn.Sequential( nn.MaxPool2d(2), DoubleConv(128, 256) ) self.down3 nn.Sequential( nn.MaxPool2d(2), DoubleConv(256, 512) ) self.down4 nn.Sequential( nn.MaxPool2d(2), DoubleConv(512, 1024) ) self.up1 nn.ConvTranspose2d(1024, 512, kernel_size2, stride2) self.conv1 DoubleConv(1024, 512) self.up2 nn.ConvTranspose2d(512, 256, kernel_size2, stride2) self.conv2 DoubleConv(512, 256) self.up3 nn.ConvTranspose2d(256, 128, kernel_size2, stride2) self.conv3 DoubleConv(256, 128) self.up4 nn.ConvTranspose2d(128, 64, kernel_size2, stride2) self.conv4 DoubleConv(128, 64) self.outc nn.Conv2d(64, n_classes, kernel_size1) def forward(self, x): x1 self.inc(x) x2 self.down1(x1) x3 self.down2(x2) x4 self.down3(x3) x5 self.down4(x4) x self.up1(x5) x torch.cat([x, x4], dim1) x self.conv1(x) x self.up2(x) x torch.cat([x, x3], dim1) x self.conv2(x) x self.up3(x) x torch.cat([x, x2], dim1) x self.conv3(x) x self.up4(x) x torch.cat([x, x1], dim1) x self.conv4(x) logits self.outc(x) return logits运行项目确保你的数据集已经放在相应的文件夹中。在项目根目录下运行以下命令启动训练bash深色版本python src/train.py训练完成后运行以下命令进行评估和可视化bash深色版本python src/predict.py功能说明数据集类RoadCrackDataset类用于加载和预处理数据。数据加载器get_data_loaders函数用于创建训练和验证数据加载器。训练模型train.py脚本用于训练UNet模型使用交叉熵损失函数和Adam优化器。评估模型predict.py脚本用于评估模型性能并可视化输入图像、真实标签和预测结果。模型定义unet.py文件定义了UNet模型。详细注释dataset.py数据集类定义了一个RoadCrackDataset类用于加载和预处理数据。数据加载器定义了一个get_data_loaders函数用于创建训练和验证数据加载器。train.py训练函数定义了一个train_model函数用于训练UNet模型。训练过程在每个epoch中模型在训练集上进行前向传播和反向传播并在验证集上进行评估。predict.py预测和可视化定义了一个predict_and_plot函数用于在验证集上进行预测并可视化输入图像、真实标签和预测结果。unet.pyUNet模型定义了UNet模型的结构包括编码器和解码器部分。