Algorithm-Repo/SourceCode/MachineLearningUtils/SemanticSeg/PytorchFramework/predict.py

# -*- coding: utf-8 -*-
# @ File       : predict.py
# @ Author     : Guido    LuXiaohao
# @ Date       : 2021/8/30
# @ Software   : PyCharm
# @ Description: 用于模型推理

import json
import math
import os
import sys
import time
from argparse import ArgumentParser
from pathlib import Path

import cv2
import numpy as np
import onnxruntime as ort
import torch
from PIL import Image, ImageDraw, ImageFont

from dataset import LoadImages, resize_LongestMaxSize
from evaluation.utils.metric_file_generate import \
    prfile_generate_semantic_segmentation

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # project root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative path

# {label: [min_area, area_rate, max_count], ...}
NECKORGAN_CFGS = {
    "甲状腺横切": [0, 1 / 25, False],
    "甲状腺纵切": [0, 1 / 25, True],
    "颈动脉短轴": [0, 1 / 100, False],
    "颈动脉长轴": [0, 1 / 100, True],
    "颈部气管": [0, 1 / 25, True]
}
NECKORGAN_CFGS2 = [
    {
        "甲状腺横切": [0, 1 / 25, False],
        "甲状腺纵切": [0, 1 / 25, True],
        "颈动脉短轴": [0, 1 / 100, False],
        "颈动脉长轴": [0, 1 / 100, True]},
    {
        "颈部气管": [0, 1 / 25, True]},
]

NECKLESION_CFGS = {
    "斑块或内中膜增厚": [50, 0, False]
}

NERVE_CFGS = {
    "神经": [0, 1 / 100, False],
    "动脉": [0, 1 / 100, False],
    "静脉": [0, 1 / 100, False]}


def resize_LongestMaxSize_back(image, ori_size, resize_mode=cv2.INTER_CUBIC):
    # 将经过resize_LongestMaxSize的图像去除pad填充，再resize回原始尺寸
    image_height = ori_size[0]
    image_width = ori_size[1]
    norm_size = image.shape[0]
    if len(image.shape) != 3:
        image = np.expand_dims(image, -1)
    if image_height > image_width:
        pad_len = int((norm_size - norm_size * image_width / image_height) / 2)
        image_before_pad = image[0:norm_size, pad_len:norm_size - pad_len]
        image = cv2.resize(image_before_pad, (image_width, image_height), interpolation=resize_mode)
    else:
        pad_len = int((norm_size - norm_size * image_height / image_width) / 2)
        image_before_pad = image[pad_len:norm_size - pad_len, 0:norm_size]
        image = cv2.resize(image_before_pad, (image_width, image_height), interpolation=resize_mode)

    return image


def semantics_segmentation_postprocess(input_tensor,
                                       label_index,
                                       min_area,
                                       area_rate,
                                       max_count=True):
    # 筛选语义分割结果
    select_mask = (input_tensor[..., label_index].squeeze() > 0.5) * 255
    select_mask = np.array(select_mask, np.uint8)
    all_area = select_mask.shape[0] * select_mask.shape[1]

    output_mask = np.zeros(select_mask.shape[:2], dtype=np.uint8)
    # 对应轮廓寻找
    contours, _ = cv2.findContours(select_mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    contours = list(contours)
    output_contours = []
    if len(contours) != 0:
        contours.sort(key=lambda x: cv2.contourArea(x), reverse=True)
        for contour_idx in range(len(contours)):
            # 轮廓为一点时，是个向量
            if np.squeeze(contours[contour_idx]).shape[0] > 10 and \
                    cv2.contourArea(contours[contour_idx]) > min_area and \
                    cv2.contourArea(contours[contour_idx]) > (all_area * area_rate):
                output_contours.append(contours[contour_idx])

    if len(output_contours) != 0:
        if max_count:
            output_mask = cv2.drawContours(output_mask.copy(), [output_contours[0]], -1, label_index, cv2.FILLED)
        else:
            output_mask = cv2.drawContours(output_mask.copy(), output_contours, -1, label_index, cv2.FILLED)

    return output_mask


def contour_sort(mask, output_label, label_mapper):
    each_image_label_list = []
    # 对应轮廓寻找
    contours, _ = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    contours = list(contours)
    contours.sort(key=lambda x: cv2.contourArea(x), reverse=True)
    for i in range(len(contours)):
        each_label_dict = {
            "contour": contours[i],
            "label_value": label_mapper[output_label],
            "area": cv2.contourArea(contours[i])}
        each_image_label_list.append(each_label_dict)

    return each_image_label_list


def draw_contours(im0, pr_contour, class_name, fontText, color_list):
    im0 = Image.fromarray(im0)  # BGR
    color = tuple([int(x) for x in color_list[pr_contour['label_value'] - 1]])  # B, G, R
    # color = (0, 0, 255)  # red
    draw = ImageDraw.Draw(im0)
    draw.text(pr_contour['contour'][0][0], class_name[pr_contour['label_value'] - 1],
              fill=color, font=fontText, stroke_width=1)
    im0 = np.asarray(im0)
    cv2.drawContours(im0, pr_contour['contour'], contourIdx=-1, color=color, thickness=2)
    return im0


def get_bounding_box(contours, img_size, expand_pixel=40):
    # 默认外扩40个像素
    roi_points = [list(i[0]) for i in contours['contour']]
    x, y = [], []
    point_count = len(roi_points)
    for ni in range(point_count):
        point_x, point_y = roi_points[ni][:]
        x.append(int(point_x))
        y.append(int(point_y))
    left, right = max(min(x) - expand_pixel, 0), min(max(x) + expand_pixel, img_size[1])  # 防止超出边界
    top, bottom = max(min(y) - expand_pixel, 0), min(max(y) + expand_pixel, img_size[0])
    return [top, bottom, left, right]


def generate_pred_file(predict_info, index_class_map):
    predict_info.sort(key=lambda x: x["Label"], reverse=False)
    output_info = [{
        "FileResultInfos": [{
            "Index": 0,
            "FrameStatus": None,
            "LabeledResult": {"Rois": []}}]}]

    for label_idx, label_info in enumerate(predict_info):
        if len(label_info["Contours"]) > 0:
            pts = [{"X": x, "Y": y} for x, y in label_info["Contours"][0]]
            output_info[0]['FileResultInfos'][0]['LabeledResult']['Rois'].append({
                "Index": label_idx,
                "Points": pts,
                "Conclusion": {
                    "Title": index_class_map[label_info["Label"]],
                    "Confidence": label_info["Confidence"]}})
        else:
            continue

    return json.dumps(output_info, ensure_ascii=False)


def model_prepare(model, num_classes, device='cuda'):
    if model.endswith(".onnx"):
        _backend = "ONNX"
    elif model.endswith((".pt", ".pth")):
        _backend = "PyTorch"
    else:
        _backend = None
        raise RuntimeError("args model must be supported format files: {'.onnx', '.pt', '.pth'}!")

    if _backend == "ONNX":
        sess_options = ort.SessionOptions()
        sess_options.execution_mode = ort.ExecutionMode.ORT_PARALLEL
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
        providers = ['CPUExecutionProvider']

        if device == "cuda":
            providers = ['CUDAExecutionProvider'] + providers
        elif device == "limit_cpu":
            sess_options.inter_op_num_threads = 1  # number of threads used to parallelize the execution across nodes
            sess_options.intra_op_num_threads = 1  # number of threads used to parallelize the execution within nodes

        # load ONNX model
        seg_model = ort.InferenceSession(model, sess_options, providers=providers)
        nc = []
        for output in seg_model.get_outputs():
            nc += [output.shape[1]]

    elif _backend == "PyTorch":
        print("You are using PyTorch as inference backend. Please make sure you call the correct model module!")
        import models
        seg_model = models.PPLiteSeg(
            num_classes=num_classes, backbone=models.STDC1(attn=True), arm_type="UAFM_ChAtten", apply_nonlin=True)
        csd = torch.load(model, map_location='cpu')['model'].float().state_dict()
        seg_model.load_state_dict(csd)  # load model state_dict
        seg_model.to(device).eval()
        nc = model.num_classes

    else:
        seg_model = None
        nc = None

    return seg_model, nc, _backend


def focus_seg(image,
              model,
              label_mapper,
              image_pre_size=(256, 256),
              resize_mode="normal",
              preprocess_mode=None,
              inference_backend="ONNX",
              postprocess_cfgs=None):
    """分割模型识别函数

    Args:
        image: 输入为裁切掉医院等无关信息的图
        model: 分割模型
        label_mapper: 用于对模型输出类别的映射
        image_pre_size: 图像预处理尺寸
        resize_mode: 图像缩放方式
        preprocess_mode: 图像预处理方式
        inference_backend: 推理后端类型
        postprocess_cfgs: 后处理配置

    Returns:
        tuple: tuple contains:
            output_contours: 模型预测的3维数组, 原图像尺寸
            output_mask: 形状为[N, H, W]的原尺寸掩模图像
            T0: 图片前处理耗时
            T1: 模型推理耗时
    """
    t0_0 = time.time()
    image_src = image.copy()
    # 尺寸放缩到固定大小
    if resize_mode == "normal":
        h = image_pre_size[0]
        w = image_pre_size[1]
        image = cv2.resize(image, (w, h))
    elif resize_mode == "fitLargeSizeAndPad":
        image = resize_LongestMaxSize(image, image_pre_size[0], resize_mode=cv2.INTER_CUBIC)
    else:
        image = image

    # 预测网络固定前处理部分
    if preprocess_mode == "normalization1":
        predict_image = (image / 255.0).astype("float32")
    elif preprocess_mode == "normalization2":
        predict_image = ((image / 255.0).astype("float32") - 0.5) / 2.0
    elif preprocess_mode == "decentralization":
        img_mean = np.mean(image)
        img_std = np.std(image)
        predict_image = ((image - img_mean + 10e-7) / (img_std + 10e-7)).astype("float32")
    elif preprocess_mode == "int8":
        predict_image = (image - 255).astype("int8")
    else:
        predict_image = image

    ############################
    t0_1 = time.time()
    T0 = t0_1 - t0_0
    print("图片前处理时间：%0.8f" % T0)

    t1_0 = time.time()
    # 模型预测
    predict_image = np.expand_dims(predict_image, axis=0)
    input_image = predict_image.transpose((0, 3, 1, 2))
    if inference_backend == "ONNX":
        input_name = model.get_inputs()[0].name
        pr_masks = model.run([], {input_name: input_image})
        pr_masks = [pr_masks] if not isinstance(pr_masks, list) else pr_masks
    elif inference_backend == "PyTorch":
        input_image = torch.from_numpy(input_image).to("cuda" if torch.cuda.is_available() else "cpu")
        pr_masks = model(input_image)
        pr_masks = [pr_masks] if not isinstance(pr_masks, list) else pr_masks
        pr_masks = [pr_mask.detach().cpu().numpy() for pr_mask in pr_masks]
    else:
        raise RuntimeError("Unsupported backend type! Only ONNX and PyTorch are supported.")

    t1_1 = time.time()
    T1 = t1_1 - t1_0
    print("模型预测时间：%0.8f" % T1)

    pr_masks = [pr_mask.squeeze(0).transpose((1, 2, 0)) for pr_mask in pr_masks]
    if resize_mode == "fitLargeSizeAndPad":
        pr_masks = [
            resize_LongestMaxSize_back(pr_mask, image_src.shape[0:2], resize_mode=cv2.INTER_LINEAR
                                       ) for pr_mask in pr_masks
        ]
    else:
        pr_masks = [
            cv2.resize(pr_mask, (image_src.shape[1], image_src.shape[0]), interpolation=cv2.INTER_LINEAR
                       ) for pr_mask in pr_masks
        ]

    # mask后处理
    output_contours = []

    for out_idx, mask in enumerate(pr_masks):
        for class_idx, label in enumerate(postprocess_cfgs[out_idx].keys()):
            min_area, area_rate, max_count = postprocess_cfgs[out_idx][label]
            pr_class_mask = semantics_segmentation_postprocess(mask, class_idx + 1, min_area, area_rate, max_count)
            pr_class_contour = contour_sort(pr_class_mask, class_idx + 1, label_mapper[out_idx])
            output_contours.extend(pr_class_contour)

    output_contours.sort(key=lambda x: x["area"], reverse=True)

    output_mask = np.zeros((image_src.shape[0], image_src.shape[1]), dtype=np.uint8)
    for single_contour in output_contours:
        output_mask = cv2.drawContours(output_mask.copy(), [single_contour["contour"]], -1,
                                       single_contour["label_value"], cv2.FILLED)

    return output_contours, output_mask, T0, T1


def run(model,
        source,
        imgsz=(256, 256),
        postprocess_cfgs=None,
        save_dir=None,
        device='cuda',
        save_txt=False):
    """用于非串联检测，如仅检测颈动脉
    """
    source = str(source)
    fontText = ImageFont.truetype(r"C:\Windows\Fonts\msyhl.ttc", 18, encoding="utf-8")
    if not isinstance(postprocess_cfgs, list):
        postprocess_cfgs = [postprocess_cfgs]
    class_name = []
    for cfg in postprocess_cfgs:
        class_name += list(cfg.keys())

    # prepare model
    seg_model, nc, _backend = model_prepare(model, 3, device)  # 3为模型的输出类别数量,依据实际情况修改
    if not isinstance(nc, list):
        nc = [nc]
    nc = [c - 1 for c in nc]
    label_mapper = [{i: i+sum(nc[:idx]) if idx > 0 else i for i in range(1, c+1)} for idx, c in enumerate(nc)]

    # Dataloader
    dataset = LoadImages(source)
    bs = len(dataset)
    vid_path, vid_writer = [None] * bs, [None] * bs

    # Run inference
    Time = []
    for path, im, im0, vid_cap, s in dataset:
        print(f'{s}')
        save_path = str(save_dir / Path(path).name)
        pr_contours, pr_focus_mask, T0, T1 = focus_seg(im, seg_model, label_mapper,
                                                       image_pre_size=imgsz,
                                                       resize_mode="normal",
                                                       preprocess_mode="normalization1",
                                                       inference_backend=_backend,
                                                       postprocess_cfgs=postprocess_cfgs)
        Time.append(T1)  # record inference time

        # draw contours
        for pr_contour in pr_contours:
            im0 = draw_contours(im0, pr_contour, class_name, fontText, COLOR_LIST_ORGAN)

        if dataset.mode == 'image':
            cv2.imencode(".jpg", im0, [int(cv2.IMWRITE_JPEG_QUALITY), 100])[1].tofile(save_path)
        else:  # 'video'
            if vid_path[dataset.count] != save_path:  # new video
                vid_path[dataset.count] = save_path
                if isinstance(vid_writer[dataset.count], cv2.VideoWriter):
                    vid_writer[dataset.count].release()  # release previous video writer
                if vid_cap:  # video
                    fps = vid_cap.get(cv2.CAP_PROP_FPS)
                    w = im0.shape[1]
                    h = im0.shape[0]
                else:
                    fps, w, h = 30, im.shape[1], im.shape[0]
                save_path = str(Path(save_path).with_suffix('.mp4'))
                vid_writer[dataset.count] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
            vid_writer[dataset.count].write(im0)

        if save_txt:
            pred_info = prfile_generate_semantic_segmentation(pr_contours, im.shape[:2])
            prf = generate_pred_file(pred_info, {k + 1: v for k, v in enumerate(class_name)})
            with open(Path(save_path).with_suffix(".txt"), "w", encoding="utf-8") as f:
                f.write(prf)

    # exclude warmup inference
    warmup = 5
    for _ in range(warmup):
        Time.pop(0)
    inference_speed = sum(Time) / len(Time)
    print('#################################\n'
          '平均每张图像推理速度为：{:.6f}\n'.format(inference_speed),
          '#################################\n')


def run_pipeline(organ_model,
                 lesion_model,
                 source,
                 imgsz=(256, 256),
                 postprocess_cfgs_organ=None,
                 postprocess_cfgs_lesion=None,
                 save_dir=None,
                 device='cuda'):
    '''
    用于串联检测，脏器+病灶，依据脏器的检测结果作为病灶的输入图像，例如颈动脉+斑块
    '''
    source = str(source)

    fontText = ImageFont.truetype(r"C:\Windows\Fonts\msyhl.ttc", 18, encoding="utf-8")
    class_name_organ = list(postprocess_cfgs_organ.keys())
    class_name_lesion = list(postprocess_cfgs_lesion.keys())

    # prepare model
    seg_organ_model, organ_backend = model_prepare(organ_model, 3, device)  # 3为模型的输出类别数量，依据实际情况修改
    seg_lesion_model, lesion_backend = model_prepare(lesion_model, 2, device)

    # Dataloader
    dataset = LoadImages(source)
    bs = len(dataset)  # batch size
    vid_path, vid_writer = [None] * bs, [None] * bs

    # Run inference
    Time = []
    for path, im, im0, vid_cap, s in dataset:
        print(f'{s}')
        save_path = str(save_dir / Path(path).name)
        pr_contours_organ, pr_focus_mask_organ, T0, T1 = \
            focus_seg(
                im, seg_organ_model,
                image_pre_size=imgsz,
                resize_mode="normal",
                preprocess_mode="normalization1",
                inference_backend=organ_backend,
                postprocess_cfgs=postprocess_cfgs_organ)

        for pr_contour_organ in pr_contours_organ:
            # TODO 需要进行病灶检测的脏器类别id，可能为多种脏器，待组成一个list
            if pr_contour_organ['label_value'] == 2:
                # 获取外接扩展矩形，如若不需要扩充像素，expand_pixel设置为0
                organ_bounding_box = get_bounding_box(pr_contour_organ,
                                                      im.shape[:2],
                                                      expand_pixel=40)
                im_lesion = im[organ_bounding_box[0]: organ_bounding_box[1],
                            organ_bounding_box[2]: organ_bounding_box[3]]

                # 裁切好的图像送入病灶检测模型
                pr_contours_lesion, pr_focus_mask_lesion, T0, T1 = \
                    focus_seg(
                        im_lesion, seg_lesion_model,
                        image_pre_size=imgsz,
                        resize_mode="fitLargeSizeAndPad",
                        preprocess_mode="normalization1",
                        inference_backend=lesion_backend,
                        postprocess_cfgs=postprocess_cfgs_lesion)

                # draw contours
                for pr_contour_lesion in pr_contours_lesion:
                    im_lesion = draw_contours(im_lesion, pr_contour_lesion, class_name_lesion, fontText,
                                              COLOR_LIST_LESION)
                im0 = im0.copy()
                im0[organ_bounding_box[0]: organ_bounding_box[1],
                organ_bounding_box[2]: organ_bounding_box[3]] = im_lesion
                im0 = draw_contours(im0, pr_contour_organ, class_name_organ, fontText, COLOR_LIST_ORGAN)

        Time.append(T1)  # record inference time

        if dataset.mode == 'image':
            cv2.imencode(".jpg", im0, [int(cv2.IMWRITE_JPEG_QUALITY), 100])[1].tofile(save_path)
        else:  # 'video'
            if vid_path[dataset.count] != save_path:  # new video
                vid_path[dataset.count] = save_path
                if isinstance(vid_writer[dataset.count], cv2.VideoWriter):
                    vid_writer[dataset.count].release()  # release previous video writer
                if vid_cap:  # video
                    fps = vid_cap.get(cv2.CAP_PROP_FPS)
                    w = im0.shape[1]
                    h = im0.shape[0]
                else:
                    fps, w, h = 30, im.shape[1], im.shape[0]
                save_path = str(Path(save_path).with_suffix('.mp4'))
                vid_writer[dataset.count] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
            vid_writer[dataset.count].write(im0)

    Time.pop(0)
    inference_speed = sum(Time) / len(Time)
    print('#################################\n'
          '平均每张图像推理速度为：{:.6f}\n'.format(inference_speed),
          '#################################\n')


def parse_args():
    parser = ArgumentParser()
    parser.add_argument(
        "--organ_model",
        help="organ model path",
        type=str,
        default=r"model.onnx")
    parser.add_argument(
        "--lesion_model",
        help="lesion model path",
        type=str,
        default=None)
    parser.add_argument(
        "--source",
        help="root path to predict",
        type=str,
        default=None)
    parser.add_argument(
        "--imgsz",
        help="inference size h,w",
        type=tuple,
        default=(256, 256))
    parser.add_argument(
        "--cfg_organ",
        help="postprocess configurations",
        default=NECKORGAN_CFGS2)
    parser.add_argument(
        "--cfg_lesion",
        help="postprocess configurations",
        default=NECKLESION_CFGS)
    parser.add_argument(
        "--save_dir",
        help="directory to save predicting results",
        type=str,
        default=None)
    parser.add_argument(
        "--device",
        help="select using cuda, cpu, or cpu with one core and one thread",
        type=str,
        choices=['cuda', 'cpu', 'limit_cpu'],
        default="cuda")
    parser.add_argument(
        "--save_txt",
        help="If true, save results to *.txt",
        type=bool,
        default=False)
    return parser.parse_args()


if __name__ == '__main__':
    opt = parse_args()
    if not isinstance(opt.cfg_organ, list):
        opt.cfg_organ = [opt.cfg_organ]
    organ_categories = []
    for cfg in opt.cfg_organ:
        organ_categories += list(cfg.keys())
    seed_arr_organ = np.array([range(1, 255, math.ceil(255 / len(organ_categories)))]).astype(np.uint8)
    COLOR_LIST_ORGAN = cv2.applyColorMap(seed_arr_organ, cv2.COLORMAP_RAINBOW)[0]
    seed_arr_lesion = np.array([range(1, 255, math.ceil(255 / len(opt.cfg_lesion)))]).astype(np.uint8)
    COLOR_LIST_LESION = cv2.applyColorMap(seed_arr_lesion, cv2.COLORMAP_RAINBOW)[0]

    # Directories
    save_dir = Path(str(opt.save_dir)) / 'predict'
    save_dir.mkdir(exist_ok=True)

    if opt.lesion_model is not None:
        # 串联检测，如检测颈动脉+斑块，目前仅支持串联两个模型，而串联三个模型的有待后续使用到再添加
        run_pipeline(
            opt.organ_model,
            opt.lesion_model,
            opt.source,
            imgsz=opt.imgsz,
            postprocess_cfgs_organ=opt.cfg_organ,
            postprocess_cfgs_lesion=opt.cfg_lesion,
            save_dir=save_dir,
            device=opt.device)
    else:
        # 非串联检测，即单个模型推理
        run(
            opt.organ_model,
            opt.source,
            imgsz=opt.imgsz,
            postprocess_cfgs=opt.cfg_organ,
            save_dir=save_dir,
            device=opt.device,
            save_txt=opt.save_txt)