当前位置：首页 > article >正文

在Ubuntu云服务器上使用OneFormer模型进行遥感图像水体提取，并替换为客户数据集的详细步骤

article 2025/2/10 5:20:36

以下是在Ubuntu云服务器上使用OneFormer模型进行遥感图像水体提取，并替换为客户数据集的详细步骤：

1. 数据集准备

首先，你需要确保客户的数据集格式与OneFormer模型兼容。通常，语义分割数据集需要包含图像和对应的标注文件（例如，图像分割的掩码）。数据集一般应按照以下目录结构组织：

dataset/
├── train/
│   ├── images/
│   │   ├── image_001.jpg
│   │   ├── image_002.jpg
│   │   └── ...
│   └── labels/
│       ├── label_001.png
│       ├── label_002.png
│       └── ...
└── val/
    ├── images/
    │   ├── image_011.jpg
    │   ├── image_012.jpg
    │   └── ...
    └── labels/
        ├── label_011.png
        ├── label_012.png
        └── ...

2. 配置数据集

在OneFormer项目中，你需要配置数据集以使用客户的数据集。具体步骤如下：

2.1 注册数据集

在项目代码中找到数据集注册的部分，通常在datasets目录下。你可以参考已有数据集的注册代码，为客户的数据集编写注册函数。以下是一个示例：

import os
from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.data.datasets import load_sem_seg

def register_custom_dataset():
    root_dir = "/path/to/dataset"  # 替换为实际的数据集路径
    for split in ["train", "val"]:
        image_dir = os.path.join(root_dir, split, "images")
        label_dir = os.path.join(root_dir, split, "labels")
        dataset_name = f"custom_{split}"
        DatasetCatalog.register(dataset_name, lambda s=split: load_sem_seg(label_dir, image_dir, gt_ext="png", image_ext="jpg"))
        MetadataCatalog.get(dataset_name).set(
            stuff_classes=["water", "other"],  # 根据实际情况修改类别名称
            stuff_colors=[(0, 255, 0), (255, 0, 0)],  # 类别对应的颜色
            image_root=image_dir,
            sem_seg_root=label_dir
        )

register_custom_dataset()

2.2 修改配置文件

找到OneFormer项目中的配置文件（通常是.yaml格式），修改数据集相关的配置项，指定使用你注册的数据集名称：

DATASETS:
  TRAIN: ("custom_train",)
  TEST: ("custom_val",)

3. 训练模型

在完成数据集配置后，你可以使用以下命令在云服务器上训练OneFormer模型：

python train_net.py --config-file configs/oneformer_swin_large_IN21k_384_bs16_100ep.yaml --num-gpus 1 SOLVER.IMS_PER_BATCH 2 SOLVER.BASE_LR 0.0001

--config-file：指定配置文件的路径。
--num-gpus：指定使用的GPU数量。
SOLVER.IMS_PER_BATCH：每个批次的图像数量。
SOLVER.BASE_LR：基础学习率。

4. 推理和水体提取

训练完成后，你可以使用训练好的模型对新的遥感图像进行推理，提取水体信息。以下是一个推理的示例代码：

import torch
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
import cv2

# 加载配置文件
cfg = get_cfg()
cfg.merge_from_file("configs/oneformer_swin_large_IN21k_384_bs16_100ep.yaml")
cfg.MODEL.WEIGHTS = "output/model_final.pth"  # 替换为实际的模型权重路径
cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# 创建预测器
predictor = DefaultPredictor(cfg)

# 读取遥感图像
image = cv2.imread("path/to/remote_sensing_image.jpg")

# 进行推理
outputs = predictor(image)

# 获取语义分割结果
sem_seg = outputs["sem_seg"].argmax(dim=0).cpu().numpy()

# 可视化结果
metadata = MetadataCatalog.get(cfg.DATASETS.TRAIN[0])
v = Visualizer(image[:, :, ::-1], metadata=metadata, scale=1.2)
out = v.draw_sem_seg(sem_seg)
cv2.imshow("Water Extraction Result", out.get_image()[:, :, ::-1])
cv2.waitKey(0)
cv2.destroyAllWindows()