基于 PyTorch 实现训练滑坡语义分割数据集 UNet 语义分割模型训练滑坡分割数据集

基于 PyTorch 实现训练滑坡语义分割数据集 UNet 语义分割模型训练滑坡分割数据集

数据集描述：
**超大规模滑坡分割数据集，共34604张图像，全部处理为512×512尺度，涵盖航空，无人机，卫星多种传感器类型，涵盖0.1m-10m不同分辨率影像，3.6GB，标注质量很高，且全部为原始数据，无增广。**的训练及应用。
1

1

滑坡语义分割数据集，共计 34604 张图像，每张大小为512×512，包含多种传感器来源（航空、无人机、卫星），分辨率范围广泛（0.1m - 10m），总大小约 3.6GB。标注为原始数据，无增广处理，非常适合用于训练 UNet 语义分割模型。

🧱 项目目标

• 使用UNet 模型对滑坡进行像素级语义分割
• 基于 PyTorch 实现训练流程
• 构建一套滑坡检测与可视化系统
• 支持单图推理、视频流或遥感影像批量推理

🛠️ 第一步：环境搭建

1. 安装 CUDA 驱动

确保你的显卡驱动已安装最新版本的 NVIDIA 驱动，并支持当前的 CUDA 版本（推荐 CUDA 11.8）：

nvidia-smi

2. 安装 Anaconda（如未安装）

从官网下载并安装 Anaconda

3. 创建虚拟环境

conda create--nameunet-segpython=3.9conda activate unet-seg

4. 安装依赖项

pipinstalltorch torchvision torchaudio pipinstallmatplotlib numpy opencv-python tqdm scikit-learn pipinstallalbumentations# 数据增强库（可选）pipinstalltensorboard# 可视化训练过程（可选）

📁 第二步：数据集结构整理（YOLO格式转UNet输入）

YOLO 格式通常如下：

dataset/ ├── images/ │ ├── train/ │ ├── val/ │ └── test/ ├── labels/ │ ├── train/ │ ├── val/ │ └── test/ └── dataset.yaml

但 UNet 要求的是图像 + mask 的配对形式，因此你需要将 YOLO 标注转换为灰度图 mask 图像。

✅ 将 YOLO 标注文件.txt转换为 mask 图像

你可以使用下面这个脚本将.txt文件转为对应尺寸的 mask 图像（每个像素值为 0 或 1）：

示例代码：convert_yolo_to_mask.py

importosimportcv2importnumpyasnp IMAGE_SIZE=(512,512)defyolo_to_mask(txt_path,img_size):h,w=img_size mask=np.zeros((h,w),dtype=np.uint8)withopen(txt_path,'r')asf:lines=f.readlines()forlineinlines:data=list(map(float,line.strip().split()))class_id=int(data[0])coords=np.array(data[1:]).reshape(-1,2)coords[:,0]*=w coords[:,1]*=h coords=np.round(coords).astype(np.int32)cv2.fillPoly(mask,[coords],color=1)# 类别为1returnmask# 输入输出路径image_dir="dataset/images/train"label_dir="dataset/labels/train"output_mask_dir="dataset/masks/train"os.makedirs(output_mask_dir,exist_ok=True)forlabel_fileinos.listdir(label_dir):base_name=os.path.splitext(label_file)[0]image_path=os.path.join(image_dir,base_name+".jpg")label_path=os.path.join(label_dir,label_file)ifnotos.path.exists(image_path):continuemask=yolo_to_mask(label_path,IMAGE_SIZE)cv2.imwrite(os.path.join(output_mask_dir,base_name+".png"),mask*255)

执行后你会得到：

dataset/ ├── images/ │ └── train/ ├── masks/ │ └── train/

🧪 第三步：构建 UNet 数据集和模型

1. 自定义 Dataset 类

importosfromtorch.utils.dataimportDatasetfromPILimportImageimportnumpyasnpimporttorchimporttorchvision.transformsastransformsclassLandslideDataset(Dataset):def__init__(self,image_dir,mask_dir,transform=None):self.image_dir=image_dir self.mask_dir=mask_dir self.transform=transform self.images=os.listdir(image_dir)def__len__(self):returnlen(self.images)def__getitem__(self,idx):img_path=os.path.join(self.image_dir,self.images[idx])mask_path=os.path.join(self.mask_dir,self.images[idx].replace('.jpg','.png'))image=np.array(Image.open(img_path).convert("RGB"))mask=np.array(Image.open(mask_path).convert("L"))# 灰度图# 二值化 mask（0 or 1）mask[mask>0]=1ifself.transform:augmentations=self.transform(image=image,mask=mask)image=augmentations["image"]mask=augmentations["mask"]image=transforms.ToTensor()(image)mask=torch.from_numpy(mask).long()returnimage,mask

2. 定义 UNet 模型（PyTorch官方实现简化版）

importtorch.nnasnnimporttorch.nn.functionalasFclassUNet(nn.Module):def__init__(self,in_channels=3,out_channels=1):super(UNet,self).__init__()defCBR(in_channels,out_channels):returnnn.Sequential(nn.Conv2d(in_channels,out_channels,kernel_size=3,padding=1),nn.BatchNorm2d(out_channels),nn.ReLU(inplace=True),nn.Conv2d(out_channels,out_channels,kernel_size=3,padding=1),nn.BatchNorm2d(out_channels),nn.ReLU(inplace=True))self.enc1=CBR(in_channels,64)self.enc2=CBR(64,128)self.enc3=CBR(128,256)self.enc4=CBR(256,512)self.bottleneck=CBR(512,1024)self.upconv4=nn.ConvTranspose2d(1024,512,kernel_size=2,stride=2)self.dec4=CBR(1024,512)self.upconv3=nn.ConvTranspose2d(512,256,kernel_size=2,stride=2)self.dec3=CBR(512,256)self.upconv2=nn.ConvTranspose2d(256,128,kernel_size=2,stride=2)self.dec2=CBR(256,128)self.upconv1=nn.ConvTranspose2d(128,64,kernel_size=2,stride=2)self.dec1=CBR(128,64)self.final_conv=nn.Conv2d(64,out_channels,kernel_size=1)defforward(self,x):enc1=self.enc1(x)enc2=self.enc2(F.max_pool2d(enc1,kernel_size=2,stride=2))enc3=self.enc3(F.max_pool2d(enc2,kernel_size=2,stride=2))enc4=self.enc4(F.max_pool2d(enc3,kernel_size=2,stride=2))bottleneck=self.bottleneck(F.max_pool2d(enc4,kernel_size=2,stride=2))dec4=self.upconv4(bottleneck)dec4=torch.cat((dec4,enc4),dim=1)dec4=self.dec4(dec4)dec3=self.upconv3(dec4)dec3=torch.cat((dec3,enc3),dim=1)dec3=self.dec3(dec3)dec2=self.upconv2(dec3)dec2=torch.cat((dec2,enc2),dim=1)dec2=self.dec2(dec2)dec1=self.upconv1(dec2)dec1=torch.cat((dec1,enc1),dim=1)dec1=self.dec1(dec1)returntorch.sigmoid(self.final_conv(dec1))

🏋️‍♂️ 第四步：训练模型

1. 训练脚本完整示例

fromtorch.utils.dataimportDataLoaderimporttorch.optimasoptimimporttorch.nnasnnfromtqdmimporttqdm# 设置参数EPOCHS=20BATCH_SIZE=8LEARNING_RATE=1e-4DEVICE="cuda"iftorch.cuda.is_available()else"cpu"# 加载数据集train_dataset=LandslideDataset(image_dir="dataset/images/train",mask_dir="dataset/masks/train")val_dataset=LandslideDataset(image_dir="dataset/images/val",mask_dir="dataset/masks/val")train_loader=DataLoader(train_dataset,batch_size=BATCH_SIZE,shuffle=True)val_loader=DataLoader(val_dataset,batch_size=BATCH_SIZE)# 初始化模型model=UNet(in_channels=3,out_channels=1).to(DEVICE)optimizer=optim.Adam(model.parameters(),lr=LEARNING_RATE)criterion=nn.BCELoss()# 二分类交叉熵损失# 开始训练forepochinrange(EPOCHS):model.train()loop=tqdm(train_loader,total=len(train_loader))forimages,masksinloop:images=images.to(DEVICE)masks=masks.float().unsqueeze(1).to(DEVICE)preds=model(images)loss=criterion(preds,masks)optimizer.zero_grad()loss.backward()optimizer.step()loop.set_postfix(loss=loss.item())# 验证阶段（可选）model.eval()withtorch.no_grad():val_loss=0forimages,masksinval_loader:images=images.to(DEVICE)masks=masks.float().unsqueeze(1).to(DEVICE)preds=model(images)val_loss+=criterion(preds,masks).item()print(f"Epoch{epoch+1}| Val Loss:{val_loss/len(val_loader):.4f}")# 保存模型torch.save(model.state_dict(),"unet_landslide.pth")

🧪 第五步：推理 & 构建滑坡检测系统

推理函数（单图测试）

defpredict_image(model,image_path,device="cuda"):model.eval()image=Image.open(image_path).convert("RGB")image_tensor=transforms.ToTensor()(image).unsqueeze(0).to(device)withtorch.no_grad():output=model(image_tensor)pred_mask=(output.squeeze().cpu().numpy()>0.5).astype(np.uint8)# 显示结果importmatplotlib.pyplotasplt plt.figure(figsize=(10,5))plt.subplot(1,2,1)plt.imshow(image)plt.title("Input Image")plt.subplot(1,2,2)plt.imshow(pred_mask,cmap='gray')plt.title("Predicted Mask")plt.show()

📦 第六步：部署 & 构建系统

你可以进一步构建以下功能：

📌 总结

你现在可以基于训练好的 UNet 模型构建一套完整的滑坡语义分割检测系统。该系统具备以下能力：

• 将 YOLO 格式的标注转化为 mask 图像
• 使用 PyTorch 实现 UNet 模型训练
• 进行单图、视频、遥感图像的滑坡检测
• 可扩展为 Web 服务或边缘设备部署

是否需要我为你提供完整的项目模板（含Flask接口、Docker打包、遥感图像处理模块等）？欢迎继续提问！