pytorch实现实例分割功能的代码

以下是使用 PyTorch 实现实例分割功能的代码示例：

import torch
import torchvision
from torchvision.models.detection import maskrcnn_resnet50_fpn

# 加载预训练的模型
model = maskrcnn_resnet50_fpn(pretrained=True)

# 设置模型为评估模式（不进行训练）
model.eval()

# 读取测试图像并进行预处理
img = Image.open('test.jpg')
transform = transforms.Compose([
    transforms.ToTensor(),
])
img_tensor = transform(img)
img_tensor = img_tensor.unsqueeze(0)

# 使用模型进行推理
with torch.no_grad():
    predictions = model(img_tensor)

# 获取预测结果中的分割掩模和类别标签
masks = predictions[0]['masks']
labels = predictions[0]['labels']

# 将掩模转换为 PIL 图像，并保存到本地
for i in range(masks.shape[0]):
    mask = masks[i, 0].mul(255).byte().cpu().numpy()
    img_mask = Image.fromarray(mask)
    img_mask.save(f'mask_{i}.png')

代码中使用的是预训练的 Mask R-CNN 模型，可以通过 `torchvision.models.detection.maskrcnn_resnet50_fpn` 加载。然后加载测试图像并进行预处理，将其转换为 PyTorch 张量。接着使用模型进行推理，获取预测结果中的分割掩模和类别标签。最后将掩模转换为 PIL 图像并保存到本地。