基于PyTorch的图像分割：如何计算船只的掩码区域？

完全可以实现，以下是从模型加载到掩码提取的全流程操作步骤：

1. 加载预训练分割模型 #

torchvision内置了多个成熟的语义分割模型（如DeepLabV3、FCN），直接调用预训练版本快速启动：

import torch
import torchvision
from torchvision.transforms import functional as F

# 加载预训练DeepLabV3模型（也可替换为fcn_resnet101等其他分割模型）
model = torchvision.models.segmentation.deeplabv3_resnet50(pretrained=True)
model.eval()  # 切换到推理模式

2. 图像预处理 #

将输入图像转换成模型要求的格式：

from PIL import Image

# 加载目标图像（替换为你的本地路径或已读取的图像对象）
image = Image.open("Ry8Nr.png")

# 转张量并增加batch维度
input_tensor = F.to_tensor(image)
input_batch = input_tensor.unsqueeze(0)

# 用模型预训练时的均值/标准差做归一化
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
input_batch = F.normalize(input_batch, mean=mean, std=std)

3. 推理并提取船只掩码 #

预训练模型基于COCO数据集训练，其中船只对应的类别索引是8，直接提取该类别的像素区域：

# 无梯度推理
with torch.no_grad():
    output = model(input_batch)['out'][0]  # 取出单张图的输出结果
output_predictions = output.argmax(0)  # 每个像素取概率最高的类别

# 生成二进制掩码：1表示船只区域，0表示背景
ship_mask = (output_predictions == 8).byte()

4. 掩码后处理（可选） #

如果掩码存在小噪声斑点，用形态学操作清理：

import cv2
import numpy as np

# 转成OpenCV兼容的numpy格式
mask_np = ship_mask.numpy().astype(np.uint8)
# 开运算（腐蚀+膨胀）去除小噪声
kernel = np.ones((5,5), np.uint8)
cleaned_mask = cv2.morphologyEx(mask_np, cv2.MORPH_OPEN, kernel)

5. 掩码可视化/保存 #

查看或保存最终的船只掩码：

import matplotlib.pyplot as plt

# 可视化掩码
plt.imshow(cleaned_mask, cmap='gray')
plt.axis('off')
plt.show()

# 保存掩码为图像文件
mask_image = Image.fromarray(cleaned_mask * 255)  # 转成0-255灰度值范围
mask_image.save("ship_mask.png")

进阶优化 #

如果预训练模型的船只分割效果不够理想，可以收集船只专属数据集（带掩码标注），对模型进行微调：

• 自定义Dataset类加载图像和标注掩码
• 用交叉熵损失作为损失函数
• 冻结模型骨干网络或全量训练，提升专属场景的分割精度

内容的提问来源于stack exchange，提问作者AHR99