两种调用mmseg中模型的简单代码

河北一帆

已于 2024-01-28 16:48:57 修改

阅读量198

点赞数 3

文章标签：人工智能

于 2024-01-28 16:37:21 首次发布

本文链接：https://blog.csdn.net/qq_41816368/article/details/135897139

版权

from mmseg.apis import MMSegInferencer


inferencer = MMSegInferencer(model='pspnet_r101b-d8_4xb2-80k_cityscapes-512x1024')
name = "/home//img/IMG_1360.jpg"
# Inference
inferencer(name, show=True)

from mmengine.config import Config
from mmengine.registry import MODELS
from mmengine.registry import init_default_scope
from mmengine.runner.checkpoint import (_load_checkpoint,
                                        _load_checkpoint_to_model)

import cv2
import torch
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt


cfg = Config.fromfile("/home/mmsegmentation/configs/pspnet/pspnet_r101b-d8_4xb2-80k_cityscapes-512x1024.py")
weights = "/home/mmsegmentation/my_checkpoints/pspnet/my_train_temp.pth"
device = "cuda"

init_default_scope('mmseg')
model = MODELS.build(cfg.model)

model.cfg = cfg
checkpoint = _load_checkpoint(weights, map_location='cpu')
_load_checkpoint_to_model(model, checkpoint)
model.to(device)
model.eval()

name = "/home/img/IMG_1360.jpg"
width = 640
height = 480
img0 = cv2.imread(name)
img0 = cv2.cvtColor(img0, cv2.COLOR_BGR2RGB)
img0 = cv2.resize(img0, (width, height))

mean = np.array([0.485, 0.456, 0.406]) * 255
std = np.array([0.229, 0.224, 0.225]) * 255
# 对图像进行归一化
img = (img0 - mean) / std

img = torch.from_numpy(img).unsqueeze(0).to(torch.float32).to(device)
img = img.permute(0, 3, 1, 2)

output = model(img)

resized_tensor = F.interpolate(output, size=(height, width), mode='bilinear', align_corners=False)

max_indices = torch.argmax(resized_tensor, dim=1)
result = torch.squeeze(max_indices, dim=1)
result = result.cpu().numpy()
result = result.astype(np.uint8)
result = result.squeeze()


plt.subplot(121)
plt.imshow(img0)
plt.subplot(122)
plt.imshow(result, cmap='viridis')
plt.title('output'), plt.xticks([]), plt.yticks([])
plt.show()