predictor.py

import os
import sys
import logging

logging.getLogger().setLevel(logging.ERROR)

import numpy as np
from copy import deepcopy

sys.path.append(
    os.path.join(os.path.dirname(os.path.realpath(__file__)), "./.."))

import paddle
import paddle.nn as nn
import paddle.nn.functional as F
from paddle.inference import create_predictor, Config

from inference.ops import DistMaps3D, ScaleLayer, BatchImageNormalize3D, SigmoidForPred


class Click:
    def __init__(self, is_positive, coords, indx=None):
        if coords is None or is_positive is None:
            raise ValueError(
                "The coord is {}, is_positive is {} and one of them is None, but none of them should be."
            )
        self.coords = coords
        self.is_positive = is_positive
        self.index = None

    @property
    def coords_and_indx(self):
        return (*self.coords, )

    def copy(self, **kwargs):
        self_copy = deepcopy(self)
        for k, v in kwargs.items():
            setattr(self_copy, k, v)
        return self_copy


class BasePredictor(object):
    def __init__(self,
                 model_path,
                 param_path,
                 net_clicks_limit=None,
                 with_mask=True,
                 norm_radius=2,
                 spatial_scale=1.0,
                 device="gpu",
                 enable_mkldnn=False,
                 **kwargs):

        self.net_clicks_limit = net_clicks_limit
        self.original_image = None
        self.prev_prediction = None
        self.model_indx = 0
        self.click_models = None
        self.net_state_dict = None
        self.with_prev_mask = with_mask
        self.device = device
        self.enable_mkldnn = enable_mkldnn
        if not paddle.in_dynamic_mode():
            paddle.disable_static()
        self.normalization = BatchImageNormalize3D(
            [0.00040428873, ],
            [0.00059983705, ], )

        self.transforms = [SigmoidForPred()]  # apply sigmoid after pred

        # !! Todo Set the radius and spatial_scale here
        self.dist_maps = DistMaps3D(
            norm_radius=norm_radius,
            spatial_scale=spatial_scale,
            cpu_mode=False,
            use_disks=True)

        # init predictor config
        self.pred_cfg = Config(model_path, param_path)
        self.pred_cfg.disable_glog_info()
        self.pred_cfg.enable_memory_optim()
        self.pred_cfg.switch_ir_optim(True)

        if self.device == "cpu":
            self._init_cpu_config()
        else:
            self._init_gpu_config()

        self.predictor = create_predictor(self.pred_cfg)

    def _init_gpu_config(self):
        logging.info("Use NVIDIA GPU")
        self.pred_cfg.enable_use_gpu(100, 0)

    def _init_cpu_config(self):
        logging.info("Use x86 CPU")
        self.pred_cfg.disable_gpu()
        if self.enable_mkldnn:
            logging.info("Use MKLDNN")
            # cache 10 different shapes for mkldnn
            # self.pred_cfg.set_mkldnn_cache_capacity(10)  # cannot use on MAC
            self.pred_cfg.enable_mkldnn()
        self.pred_cfg.set_cpu_math_library_num_threads(10)

    def set_input_image(self, image):  # (1, 12, 512, 512)
        # image is np array or other format including scalar，tuple，list，and paddle.Tensor
        self.original_image = paddle.to_tensor(image).astype(
            "float32") / 255  # 1, 512, 512, 12]

        for transform in self.transforms:
            transform.reset()

        if len(self.original_image.shape) == 4:
            self.original_image = self.original_image.unsqueeze(
                0)  # (1, 1, 12, 512, 512)

        # 默认 concate 一个全0的mask作为prev mask
        self.prev_prediction = paddle.zeros_like(
            self.original_image[:, :1, :, :, :])

        if not self.with_prev_mask:
            self.prev_edge = paddle.zeros_like(self.original_image[:, :
                                                                   1, :, :, :])

    def get_prediction_noclicker(self, clicker, prev_mask=None):
        clicks_list = clicker.get_clicks()  # one click a time todo：累计多个点

        input_image = self.original_image  # [1, 1, 512, 512, 12]
        if prev_mask is None:
            if not self.with_prev_mask:
                prev_mask = self.prev_edge
            else:
                prev_mask = self.prev_prediction
        input_image = paddle.concat(
            [input_image, prev_mask], axis=1)  # [1, 2, 512, 512, 12]

        image_nd, clicks_lists, is_image_changed = self.apply_transforms(
            input_image, [clicks_list])
        pred_logits = self._get_prediction(image_nd, clicks_lists,
                                           is_image_changed)

        pred_logits = paddle.to_tensor(pred_logits)  # [1, 1, 512, 512, 12]
        for t in reversed(self.transforms):  # conv as the final output
            pred_logits = t.inv_transform(pred_logits)

        self.prev_prediction = pred_logits
        return pred_logits.numpy()[0, 0]

    def get_prediction(self, clicker, prev_mask=None):
        clicks_list = clicker.get_clicks()  #

        input_image = self.original_image
        if prev_mask is None:
            if not self.with_prev_mask:
                prev_mask = self.prev_edge
            else:
                prev_mask = self.prev_prediction

        input_image = paddle.concat([input_image, prev_mask], axis=1)

        image_nd, clicks_lists, is_image_changed = self.apply_transforms(
            input_image, [clicks_list])
        pred_logits = self._get_prediction(image_nd, clicks_lists,
                                           is_image_changed)

        pred_logits = paddle.to_tensor(pred_logits)

        prediction = F.interpolate(
            pred_logits,
            mode="trilinear",
            align_corners=True,
            size=image_nd.shape[2:],
            data_format="NCDHW")

        for t in reversed(self.transforms):
            if pred_edges is not None:
                edge_prediction = t.inv_transform(edge_prediction)
                self.prev_edge = edge_prediction
            prediction = t.inv_transform(prediction)

        self.prev_prediction = prediction
        return prediction.numpy()[0, 0]

    def prepare_input(self, image):
        prev_mask = image[:, 1:, :, :, :]
        image = image[:, :1, :, :, :]
        image = self.normalization(image)
        return image, prev_mask

    def get_coord_features(self, image, prev_mask, points):
        coord_features = self.dist_maps(image, points)  # [1, 2, 512, 512, 12]

        if prev_mask is not None:
            coord_features = paddle.concat(
                (prev_mask, coord_features), axis=1)  # [1, 3, 512, 512, 12]

        return coord_features

    def _get_prediction(
            self, image_nd, clicks_lists,
            is_image_changed):  # what is the click? click on the right place?
        input_names = self.predictor.get_input_names()
        self.input_handle_1 = self.predictor.get_input_handle(input_names[0])
        self.input_handle_2 = self.predictor.get_input_handle(input_names[1])
        points_nd = self.get_points_nd(clicks_lists)  # 一个正点，一个负点

        image, prev_mask = self.prepare_input(image_nd)

        coord_features = self.get_coord_features(image, prev_mask, points_nd)
        image = image.numpy().astype("float32")
        coord_features = coord_features.numpy().astype("float32")

        # logging.info("coord_features.shape, image.shape", coord_features.shape, image.shape, prev_mask.shape)
        self.input_handle_1.copy_from_cpu(image)
        self.input_handle_2.copy_from_cpu(coord_features)

        self.predictor.run()

        output_names = self.predictor.get_output_names()

        output_handle = self.predictor.get_output_handle(output_names[0])
        output_data = output_handle.copy_to_cpu()
        return output_data

    def _get_transform_states(self):
        return [x.get_state() for x in self.transforms]

    def _set_transform_states(self, states):
        assert len(states) == len(self.transforms)
        for state, transform in zip(states, self.transforms):
            transform.set_state(state)

    def apply_transforms(self, image_nd, clicks_lists):
        is_image_changed = False
        for t in self.transforms:
            image_nd, clicks_lists = t.transform(image_nd, clicks_lists)
            is_image_changed |= t.image_changed

        return image_nd, clicks_lists, is_image_changed

    def get_points_nd(self, clicks_lists):
        total_clicks = []
        logging.info(
            "check_list",
            clicks_lists, )
        num_pos_clicks = [
            sum(x.is_positive for x in clicks_list)
            for clicks_list in clicks_lists
        ]
        num_neg_clicks = [
            len(clicks_list) - num_pos
            for clicks_list, num_pos in zip(clicks_lists, num_pos_clicks)
        ]

        num_max_points = max(num_pos_clicks + num_neg_clicks)

        if self.net_clicks_limit is not None:
            num_max_points = min(self.net_clicks_limit, num_max_points)
        num_max_points = max(1, num_max_points)

        for clicks_list in clicks_lists:
            clicks_list = clicks_list[:self.net_clicks_limit]
            pos_clicks = [
                click.coords_and_indx for click in clicks_list
                if click.is_positive
            ]
            pos_clicks = pos_clicks + (num_max_points - len(pos_clicks)
                                       ) * [(-1, -1, -1, -1)]

            neg_clicks = [
                click.coords_and_indx for click in clicks_list
                if not click.is_positive
            ]
            neg_clicks = neg_clicks + (num_max_points - len(neg_clicks)
                                       ) * [(-1, -1, -1, -1)]
            total_clicks.append(pos_clicks + neg_clicks)

        return paddle.to_tensor(total_clicks)

    def get_states(self):
        return {
            "transform_states": self._get_transform_states(),
            "prev_prediction": self.prev_prediction,
        }

    def set_states(self, states):
        self._set_transform_states(states["transform_states"])
        self.prev_prediction = states["prev_prediction"]