手把手实战教学，语义分割从0到1：三、推理_批量推理

本篇博客，是《手把手实战教学！语义分割从0到1》系列的第三篇实战教学，将重点介绍语义分割推理部分。

本系列总的介绍，以及其他章节的汇总，见：https://blog.csdn.net/oYeZhou/article/details/115123589。

目录

1、工程上进行推理的要素

2、构造推理工程

2.1、保留必要的内容

2.2、编写推理class

2.3、单图推理

2.4、批量推理

3、写在后面

---

1、工程上进行推理的要素

不管是语义分割还是目标检测或者其他的任务，我们要想在工程上进行模型推理，一般都是这么个流程（这里不做量化加速）：

• （1）抽离模型定义及权重文件；
• （2）编写推理class；
• （3）提供单张推理及数据集推理的方法；

因此，我们需要在推理的工程中，准备这样几种东西：模型定义脚本、权重文件、推理class、单张测试脚本、批量测试（评价）脚本。

2、构造推理工程

2.1、保留必要的内容

在推理工程中，我们仅保留推理所必须的部分，其他的代码、文件统统去掉。

按照我们上篇博客的介绍，我们已经使用自己的VOC格式的语义分割数据集训练好了一个DeepLabv3+模型。此时，我们把其中对推理有必要的部分抽离出来，包括：

上图所示的部分是模型的定义。

然后，把训练好的模型单独放在一个文件夹中：

2.2、编写推理class

接下来，就是编写推理class了。在该class中，需要包含模型定义、模型权重加载、单张推理、结果可视化等内容。这里贴出我所编写的class内容：

# -*- coding: utf-8 -*-
"""
Created on 2020.01.12
@author: LWS
Inference of Edge Fence Segmentation.
"""
 
import numpy as np
from PIL import Image
from collections import OrderedDict
 
import torch
import torch.nn.functional as F
from torchvision import transforms
 
from models.deeplabv3_plus import DeepLab
from utils.helpers import colorize_mask
from utils import palette
 
 
class EdgeFenceSeg(object):
 
    def __init__(self,
                 model_path="ckpt/best_model.pth",
                 area_thres=0.1,
                 cuda_id=0,
                 get_mask=True):
 
        torch.set_num_threads(8)
 
        self.area_thres = area_thres
        self.num_classes = 2  # background + fence
        self.get_mask = get_mask
 
        # get palette of VOC
        self.my_palette = palette.get_voc_palette(self.num_classes)
 
        # data setting
        self._MEAN = [0.48311856, 0.49071315, 0.45774156]
        self._STD = [0.21628413, 0.22036915, 0.22477823]
        self.to_tensor = transforms.ToTensor()
        self.normalize = transforms.Normalize(self._MEAN, self._STD)
 
        # get Model
        self.model = DeepLab(num_classes=self.num_classes, backbone='resnet101', pretrained=False)
        availble_gpus = list(range(torch.cuda.device_count()))
        self.device = torch.device('cuda:{}'.format(cuda_id) if len(availble_gpus) > 0 else 'cpu')
 
        # Load checkpoint
        checkpoint = torch.load(model_path, map_location=self.device)
        if isinstance(checkpoint, dict) and 'state_dict' in checkpoint.keys():
            checkpoint = checkpoint['state_dict']
        # If during training, we used data parallel
        if ('module' in list(checkpoint.keys())[0] and
                not isinstance(self.model, torch.nn.DataParallel)):
            # for gpu inference, use data parallel
            if "cuda" in self.device.type:
                self.model = torch.nn.DataParallel(self.model)
            else:
                # for cpu inference, remove module
                new_state_dict = OrderedDict()
                for k, v in checkpoint.items():
                    name = k[7:]
                    new_state_dict[name] = v
                checkpoint = new_state_dict
        # load
        self.model.load_state_dict(checkpoint)
        self.model.to(self.device)
        self.model.eval()
 
    def predict(self, img):
        """
        :param img: image for predict, np.ndarray.
        :return: mask_img, prediction, flag;
        if all None, means image type error; if mask_img is None, means don't extract mask.
        """
        if str(type(img)) == "<class 'NoneType'>":
            return None, None, None
        flag = False
        if isinstance(img, np.ndarray):
            img = Image.fromarray(img)
        with torch.no_grad():
            input = self.normalize(self.to_tensor(img)).unsqueeze(0)
            prediction = self.model(input.to(self.device))
            prediction = prediction.squeeze(0).cpu().numpy()
            prediction = F.softmax(torch.from_numpy(prediction),
                                   dim=0).argmax(0).cpu().numpy()
            area_ratio = sum(prediction[prediction == 1])/(img.size[0]*img.size[1])
            if area_ratio >= self.area_thres:
                flag = True
        if self.get_mask:
            mask_img = self.colored_mask_img(img, prediction)
            return mask_img, prediction, flag
        else:
            return None, prediction, flag
 
    def colored_mask_img(self, image, mask):
        colorized_mask = colorize_mask(mask, self.my_palette)
        # PIL type
        mask_img = Image.blend(image.convert('RGBA'), colorized_mask.convert('RGBA'), 0.7)
        return mask_img
 

2.3、单图推理

可以定义一个main.py文件，利用上述class进行单张图片的推理：

import os
import time
import cv2
import numpy as np
 
from EdgeFenceSeg import EdgeFenceSeg
 
if __name__ == "__main__":
    img_file = "test_imgs/V10108-115508_frame_232.jpg"
    output_path = "output_cv"
    if not os.path.exists(output_path):
        os.makedirs(output_path)
 
    edg = EdgeFenceSeg(area_thres=0.1, cuda_id=0, get_mask=True)
    img = cv2.imread(img_file)
 
    for i in range(4):
        # inference
        t1 = time.time()
        mask_img, prediction, flag = edg.predict(img)
        t2 = time.time()
        print("time: {}, is edge_fence: {}".format(round(t2 - t1, 4), flag))
 
    # save masked img
    if mask_img is not None:
        image_file = os.path.basename(img_file).split('.')[0]
        # mask_img_cv = cv2.cvtColor(np.asarray(mask_img), cv2.COLOR_RGB2BGR)
        mask_img_cv = np.asarray(mask_img)
        cv2.imwrite(os.path.join(output_path, image_file + '.png'), mask_img_cv)

2.4、批量推理

如果想对一批图片进行推理，由于我们的网络是全卷积的，所以输入图片的尺寸可以是任意的，这里我们使用原图大小以获取更好的性能，因此批量推理也是加个for循环进行的逐张推理：

import os
import time
from glob import glob
import cv2
import numpy as np
 
from EdgeFenceSeg import EdgeFenceSeg
 
if __name__ == "__main__":
    imgs_path = "test_imgs"
    output_path = "output_cv"
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    edg = EdgeFenceSeg()
    image_files = sorted(glob(os.path.join(imgs_path, f'*.{"jpg"}')))
    for img_file in image_files:
        t0 = time.time()
        img = cv2.imread(img_file)
 
        # inference
        t1 = time.time()
        mask_img, prediction, flag = edg.predict(img)
        t2 = time.time()
        print("{0:50}: Inference time: {1},   Is edge_fence: {2}".format(img_file, round(t2 - t1, 4), flag))
 
        # save masked img
        if mask_img is not None:
            image_file = os.path.basename(img_file).split('.')[0]
            # mask_img_cv = cv2.cvtColor(np.asarray(mask_img), cv2.COLOR_RGB2BGR)
            mask_img_cv = np.asarray(mask_img)
            cv2.imwrite(os.path.join(output_path, image_file + '.png'), mask_img_cv)

3、写在后面

本篇博客是本系列实战教程的最后一篇了，至此，如果你按照这几篇博客的描述一步步走过来，应该已经掌握了语义分割的入坑基本流程了。