使用 U-Net 进行语义分割任务 无人机树冠类分割数据集的训练 分割无人机树冠 使用 **PyTorch torchvision COCO API文章目录使用 U-Net 进行语义分割任务 无人机树冠类分割数据集的训练 分割无人机树冠 使用 **PyTorch torchvision COCO API**无人机树冠类分割数据集****数据集划分****标注信息COCO格式**✅ 三、将 COCO JSON 转为二值 Mask 图像utils.py✅ 四、U-Net 模型定义unet_model.py✅ 五、自定义 Datasetdataset.py✅ 六、训练脚本train.py✅ 七、推理与可视化infer.py✅ 九、总结数据集文字描述4633 张无人机航拍视角的图像 无人机树冠类分割数据集数据集划分为两个部分Train集: 3243 张图片Val集: 1390 张图片已完成分割标注为json的COCO格式111无人机树冠类分割数据集项目内容数据集名称无人机树冠类分割数据集图像总数4,633 张数据来源无人机航拍视角任务类型实例分割Instance Segmentation标注格式JSON 格式COCO 标准格式标注内容每棵树冠的多边形轮廓polygon分割掩码数据集划分划分类型图像数量占比用途训练集 (Train)3,243 张~70%模型训练验证集 (Val)1,390 张~30%模型验证与调优标注信息COCO格式字段说明images包含每张图像的文件名、宽度、高度、ID等信息annotations每个实例的分割多边形坐标segmentation、边界框bbox、类别IDcategory_id、图像IDimage_id等categories类别定义例如{ id: 1, name: tree, supercategory: plant }文件结构示例annotations_train.jsonannotations_val.jsontrain/├── image_0001.jpg├── image_0002.jpg└── …val/├── image_3244.jpg└── …--- ### **数据集** - ✅ **高质量航拍图像**来自无人机视角分辨率高适合林业监测 - ✅ **完整实例分割标注**每棵树冠均以多边形标注可用于精确分割任务 - ✅ **标准COCO格式**兼容主流框架如 Detectron2、MMDetection、YOLOv8-seg - ✅ **合理划分**训练集与验证集已划分好便于直接训练与评估 --- 使用 **U-Net** 进行语义分割任务。使用 **PyTorch torchvision COCO API** 完成以下流程 - ✅ 数据预处理COCO JSON → Mask 图像 - ✅ 构建 U-Net 模型 - ✅ 数据加载器 - ✅ 训练与验证代码 - ✅ 推理与可视化 --- ## 一、项目结构建议tree_canopy_segmentation/├── dataset/│ ├── images/│ │ ├── train/│ │ └── val/│ ├── annotations/│ │ ├── train.json│ │ └── val.json│ └── masks/ # 自动生成的二值 mask│ ├── train/│ └── val/├── unet_model.py # U-Net 定义├── dataset.py # 自定义 Dataset├── train.py # 训练脚本├── infer.py # 推理脚本├── utils.py # 工具函数COCO转mask└── requirements.txt--- ## ✅ 二、环境安装 bash # 创建虚拟环境 conda create -n unet_tree python3.9 -y conda activate unet_tree # 安装依赖 pip install torch torchvision opencv-python matplotlib pycocotools scikit-image pillow✅ 三、将 COCO JSON 转为二值 Mask 图像utils.py# utils.pyimportosfrompycocotools.cocoimportCOCOfromPILimportImageimportnumpyasnpimportcv2defcoco_to_masks(coco_json,image_dir,mask_dir): 将 COCO 格式标注转换为单类二值 mask 图像 :param coco_json: COCO 格式 json 文件路径 :param image_dir: 原图路径 :param mask_dir: 输出 mask 路径 os.makedirs(mask_dir,exist_okTrue)cocoCOCO(coco_json)cat_idscoco.getCatIds()# 获取所有类别 IDimg_idscoco.getImgIds()forimg_idinimg_ids:img_infococo.loadImgs(img_id)[0]ann_idscoco.getAnnIds(imgIdsimg_id,catIdscat_ids,iscrowdNone)annscoco.loadAnns(ann_ids)# 读取图像大小h,wimg_info[height],img_info[width]masknp.zeros((h,w),dtypenp.uint8)# 合并所有实例的 maskforanninanns:seg_maskcoco.annToMask(ann)masknp.maximum(mask,seg_mask)# 保存为 0/255 二值图mask(mask*255).astype(np.uint8)mask_pathos.path.join(mask_dir,img_info[file_name].replace(.jpg,_mask.png))Image.fromarray(mask).save(mask_path)print(fMask 生成完成保存至{mask_dir})# 转换训练集和验证集if__name____main__:coco_to_masks(coco_jsondataset/annotations/train.json,image_dirdataset/images/train,mask_dirdataset/masks/train)coco_to_masks(coco_jsondataset/annotations/val.json,image_dirdataset/images/val,mask_dirdataset/masks/val) 运行一次即可生成masks/train/和masks/val/中的二值 mask 图像。✅ 四、U-Net 模型定义unet_model.py# unet_model.pyimporttorchimporttorch.nnasnndefdouble_conv(in_channels,out_channels):returnnn.Sequential(nn.Conv2d(in_channels,out_channels,3,padding1),nn.BatchNorm2d(out_channels),nn.ReLU(inplaceTrue),nn.Conv2d(out_channels,out_channels,3,padding1),nn.BatchNorm2d(out_channels),nn.ReLU(inplaceTrue))classUNet(nn.Module):def__init__(self,n_classes1):super(UNet,self).__init__()self.dconv_down1double_conv(3,64)self.dconv_down2double_conv(64,128)self.dconv_down3double_conv(128,256)self.dconv_down4double_conv(256,512)self.maxpoolnn.MaxPool2d(2)self.upsamplenn.Upsample(scale_factor2,modebilinear,align_cornersTrue)self.dconv_up3double_conv(256512,256)self.dconv_up2double_conv(128256,128)self.dconv_up1double_conv(12864,64)self.conv_lastnn.Conv2d(64,n_classes,1)self.sigmoidnn.Sigmoid()ifn_classes1elsenn.Identity()defforward(self,x):# 下采样conv1self.dconv_down1(x)xself.maxpool(conv1)conv2self.dconv_down2(x)xself.maxpool(conv2)conv3self.dconv_down3(x)xself.maxpool(conv3)xself.dconv_down4(x)# 上采样 跳跃连接xself.upsample(x)xtorch.cat([x,conv3],dim1)xself.dconv_up3(x)xself.upsample(x)xtorch.cat([x,conv2],dim1)xself.dconv_up2(x)xself.upsample(x)xtorch.cat([x,conv1],dim1)xself.dconv_up1(x)xself.conv_last(x)ifself.sigmoidisnotNone:xself.sigmoid(x)returnx✅ 五、自定义 Datasetdataset.py# dataset.pyimportosfromtorch.utils.dataimportDatasetfromPILimportImageimporttorchvision.transformsasTclassTreeCanopyDataset(Dataset):def__init__(self,image_dir,mask_dir,transformNone):self.image_dirimage_dir self.mask_dirmask_dir self.transformtransform self.imagessorted(os.listdir(image_dir))self.maskssorted(os.listdir(mask_dir))def__len__(self):returnlen(self.images)def__getitem__(self,idx):img_pathos.path.join(self.image_dir,self.images[idx])mask_pathos.path.join(self.mask_dir,self.masks[idx])imageImage.open(img_path).convert(RGB)maskImage.open(mask_path).convert(L)# 灰度图ifself.transform:imageself.transform(image)maskT.ToTensor()(mask)# [1, H, W]mask(mask0.5).float()# 二值化returnimage,mask✅ 六、训练脚本train.py# train.pyimporttorchimporttorch.nnasnnfromtorch.utils.dataimportDataLoaderfromunet_modelimportUNetfromdatasetimportTreeCanopyDatasetimporttorchvision.transformsasTfromtqdmimporttqdm# 参数设置DEVICEcudaiftorch.cuda.is_available()elsecpuEPOCHS100LR1e-4BATCH_SIZE8IMG_SIZE256# 数据增强transformT.Compose([T.Resize((IMG_SIZE,IMG_SIZE)),T.ToTensor(),T.Normalize(mean[0.485,0.456,0.406],std[0.229,0.224,0.225])])# 数据集train_datasetTreeCanopyDataset(image_dirdataset/images/train,mask_dirdataset/masks/train,transformtransform)val_datasetTreeCanopyDataset(image_dirdataset/images/val,mask_dirdataset/masks/val,transformtransform)train_loaderDataLoader(train_dataset,batch_sizeBATCH_SIZE,shuffleTrue)val_loaderDataLoader(val_dataset,batch_sizeBATCH_SIZE,shuffleFalse)# 模型modelUNet(n_classes1).to(DEVICE)criterionnn.BCELoss()# 二分类交叉熵optimizertorch.optim.Adam(model.parameters(),lrLR)# 训练循环best_lossfloat(inf)forepochinrange(EPOCHS):model.train()train_loss0.0forimages,masksintqdm(train_loader,descfEpoch{epoch1}/{EPOCHS}):images,masksimages.to(DEVICE),masks.to(DEVICE)optimizer.zero_grad()outputsmodel(images)losscriterion(outputs,masks)loss.backward()optimizer.step()train_lossloss.item()# 验证model.eval()val_loss0.0withtorch.no_grad():forimages,masksinval_loader:images,masksimages.to(DEVICE),masks.to(DEVICE)outputsmodel(images)losscriterion(outputs,masks)val_lossloss.item()train_loss/len(train_loader)val_loss/len(val_loader)print(fEpoch [{epoch1}/{EPOCHS}], Train Loss:{train_loss:.4f}, Val Loss:{val_loss:.4f})# 保存最佳模型ifval_lossbest_loss:best_lossval_loss torch.save(model.state_dict(),weights/best_unet_tree.pth)print(f模型已保存: weights/best_unet_tree.pth)✅ 七、推理与可视化infer.py# infer.pyimporttorchfromPILimportImageimportnumpyasnpimportmatplotlib.pyplotaspltfromunet_modelimportUNetfromtorchvisionimporttransforms# 加载模型modelUNet(n_classes1)model.load_state_dict(torch.load(weights/best_unet_tree.pth,map_locationcpu))model.eval()# 预处理transformtransforms.Compose([transforms.Resize((256,256)),transforms.ToTensor(),transforms.Normalize(mean[0.485,0.456,0.406],std[0.229,0.224,0.225])])# 推理单张图像defpredict(image_path):imageImage.open(image_path).convert(RGB)orig_sizeimage.size# (W, H)input_tensortransform(image).unsqueeze(0)# [1, 3, 256, 256]withtorch.no_grad():outputmodel(input_tensor)outputtorch.sigmoid(output)output(output0.5).float().squeeze().cpu().numpy()# [H, W]# 可视化plt.figure(figsize(12,6))plt.subplot(1,2,1)plt.imshow(image)plt.title(Original Image)plt.axis(off)plt.subplot(1,2,2)plt.imshow(output,cmapgray)plt.title(Predicted Mask)plt.axis(off)plt.show()# 使用predict(dataset/images/val/tree_001.jpg)✅ 九、总结你已具备完整流程✅ COCO JSON → 二值 Mask✅ U-Net 模型定义✅ 数据加载与训练✅ 推理可视化运行命令python utils.py# 生成 maskpython train.py# 开始训练python infer.py# 测试推理如需扩展支持多类树种分割