<reponame>bpanahij/maskal
# @Author: <NAME>
# @Date:   2021-03-22 09:43:07
# @Last Modified by:   <NAME>
# @Last Modified time: 2021-07-12 13:33:36

from itertools import chain 

## based on: detectron2.engine.defaults
import torch
import detectron2.data.transforms as T
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.data import MetadataCatalog
from detectron2.modeling import build_model
from baal.bayesian.dropout import patch_module

__all__ = ["MonteCarloDropout", "MonteCarloDropoutHead"]


class MonteCarloDropout:
    """
    Create a simple end-to-end predictor with the given config that runs on
    single device for a single input image.

    Compared to using the model directly, this class does the following additions:

    1. Load checkpoint from `cfg.MODEL.WEIGHTS`.
    2. Always take BGR image as the input and apply conversion defined by `cfg.INPUT.FORMAT`.
    3. Apply resizing defined by `cfg.INPUT.{MIN,MAX}_SIZE_TEST`.
    4. Take one input image and produce a single output, instead of a batch.

    If you'd like to do anything more fancy, please refer to its source code
    as examples to build and use the model manually.

    Attributes:
        metadata (Metadata): the metadata of the underlying dataset, obtained from
            cfg.DATASETS.TEST.

    Examples:
    ::
        pred = DefaultPredictor(cfg)
        inputs = cv2.imread("input.jpg")
        outputs = pred(inputs)
    """

    def __init__(self, cfg, iterations=1, batch_size=1):
        self.cfg = cfg.clone()  # cfg can be modified by model
        self.model = build_model(self.cfg)
        self.model = patch_module(self.model)
        self.model.eval()
        if len(cfg.DATASETS.TEST):
            self.metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])

        checkpointer = DetectionCheckpointer(self.model)
        checkpointer.load(cfg.MODEL.WEIGHTS)

        self.aug = T.ResizeShortestEdge(
            [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST
        )

        self.input_format = cfg.INPUT.FORMAT
        assert self.input_format in ["RGB", "BGR"], self.input_format

        self.iterations = iterations
        self.batch_size = batch_size

    def __call__(self, original_image):
        """
        Args:
            original_image (np.ndarray): an image of shape (H, W, C) (in BGR order).

        Returns:
            predictions (dict):
                the output of the model for one image only.
                See :doc:`/tutorials/models` for details about the format.
        """
        with torch.no_grad():  # https://github.com/sphinx-doc/sphinx/issues/4258
            # Apply pre-processing to image.
            if self.input_format == "RGB":
                # whether the model expects BGR inputs or RGB
                original_image = original_image[:, :, ::-1]
            height, width = original_image.shape[:2]
            image = self.aug.get_transform(original_image).apply_image(original_image)
            image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))

            inputs = []
            for _ in range(self.iterations):
                inputs.append({"image": image, "height": height, "width": width})

            batch_inputs = [inputs[i:i + self.batch_size] for i in range(0, self.iterations, self.batch_size)]  

            prediction_list = []
            for b in range(len(batch_inputs)):
                prediction_list.append(self.model(batch_inputs[b]))
            predictions = list(chain.from_iterable(prediction_list)) 

            return predictions


class MonteCarloDropoutHead:
    """
    Create a simple end-to-end predictor with the given config that runs on
    single device for a single input image.

    Compared to using the model directly, this class does the following additions:

    1. Load checkpoint from `cfg.MODEL.WEIGHTS`.
    2. Always take BGR image as the input and apply conversion defined by `cfg.INPUT.FORMAT`.
    3. Apply resizing defined by `cfg.INPUT.{MIN,MAX}_SIZE_TEST`.
    4. Take one input image and produce a single output, instead of a batch.

    If you'd like to do anything more fancy, please refer to its source code
    as examples to build and use the model manually.

    Attributes:
        metadata (Metadata): the metadata of the underlying dataset, obtained from
            cfg.DATASETS.TEST.

    Examples:
    ::
        pred = DefaultPredictor(cfg)
        inputs = cv2.imread("input.jpg")
        outputs = pred(inputs)
    """

    def __init__(self, cfg, iterations=1):
        self.cfg = cfg.clone()  # cfg can be modified by model
        self.model = build_model(self.cfg)
        self.model = patch_module(self.model)
        self.model.eval()
        if len(cfg.DATASETS.TEST):
            self.metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])

        checkpointer = DetectionCheckpointer(self.model)
        checkpointer.load(cfg.MODEL.WEIGHTS)

        self.aug = T.ResizeShortestEdge(
            [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST
        )

        self.input_format = cfg.INPUT.FORMAT
        assert self.input_format in ["RGB", "BGR"], self.input_format

        self.iterations = iterations

    def __call__(self, original_image):
        """
        Args:
            original_image (np.ndarray): an image of shape (H, W, C) (in BGR order).

        Returns:
            predictions (dict):
                the output of the model for one image only.
                See :doc:`/tutorials/models` for details about the format.
        """
        with torch.no_grad():  # https://github.com/sphinx-doc/sphinx/issues/4258
            # Apply pre-processing to image.
            if self.input_format == "RGB":
                # whether the model expects BGR inputs or RGB
                original_image = original_image[:, :, ::-1]
            height, width = original_image.shape[:2]
            image = self.aug.get_transform(original_image).apply_image(original_image)
            image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
            inputs = [{"image": image, "height": height, "width": width}]

            images = self.model.preprocess_image(inputs)
            features = self.model.backbone(images.tensor)
            proposals, _ = self.model.proposal_generator(images, features, None)

            features_ = [features[f] for f in self.model.roi_heads.box_in_features]
            box_features_pooler = self.model.roi_heads.box_pooler(features_, [x.proposal_boxes for x in proposals])
            
            prediction_list = []
            for _ in range(self.iterations):
                box_features = self.model.roi_heads.box_head(box_features_pooler)
                predictions = self.model.roi_heads.box_predictor(box_features)
                pred_instances, pred_inds = self.model.roi_heads.box_predictor.inference(predictions, proposals)
                pred_instances = self.model.roi_heads.forward_with_given_boxes(features, pred_instances)
                outputs = self.model._postprocess(pred_instances, inputs, images.image_sizes)
                prediction_list.append(outputs)
            predictions = list(chain.from_iterable(prediction_list)) 

            return predictions
# @Author: <NAME>
# @Date:   2021-03-22 09:43:07
# @Last Modified by:   <NAME>
# @Last Modified time: 2021-07-12 13:33:36

from itertools import chain 

## based on: detectron2.engine.defaults
import torch
import detectron2.data.transforms as T
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.data import MetadataCatalog
from detectron2.modeling import build_model
from baal.bayesian.dropout import patch_module

__all__ = ["MonteCarloDropout", "MonteCarloDropoutHead"]


class MonteCarloDropout:
    """
    Create a simple end-to-end predictor with the given config that runs on
    single device for a single input image.

    Compared to using the model directly, this class does the following additions:

    1. Load checkpoint from `cfg.MODEL.WEIGHTS`.
    2. Always take BGR image as the input and apply conversion defined by `cfg.INPUT.FORMAT`.
    3. Apply resizing defined by `cfg.INPUT.{MIN,MAX}_SIZE_TEST`.
    4. Take one input image and produce a single output, instead of a batch.

    If you'd like to do anything more fancy, please refer to its source code
    as examples to build and use the model manually.

    Attributes:
        metadata (Metadata): the metadata of the underlying dataset, obtained from
            cfg.DATASETS.TEST.

    Examples:
    ::
        pred = DefaultPredictor(cfg)
        inputs = cv2.imread("input.jpg")
        outputs = pred(inputs)
    """

    def __init__(self, cfg, iterations=1, batch_size=1):
        self.cfg = cfg.clone()  # cfg can be modified by model
        self.model = build_model(self.cfg)
        self.model = patch_module(self.model)
        self.model.eval()
        if len(cfg.DATASETS.TEST):
            self.metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])

        checkpointer = DetectionCheckpointer(self.model)
        checkpointer.load(cfg.MODEL.WEIGHTS)

        self.aug = T.ResizeShortestEdge(
            [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST
        )

        self.input_format = cfg.INPUT.FORMAT
        assert self.input_format in ["RGB", "BGR"], self.input_format

        self.iterations = iterations
        self.batch_size = batch_size

    def __call__(self, original_image):
        """
        Args:
            original_image (np.ndarray): an image of shape (H, W, C) (in BGR order).

        Returns:
            predictions (dict):
                the output of the model for one image only.
                See :doc:`/tutorials/models` for details about the format.
        """
        with torch.no_grad():  # https://github.com/sphinx-doc/sphinx/issues/4258
            # Apply pre-processing to image.
            if self.input_format == "RGB":
                # whether the model expects BGR inputs or RGB
                original_image = original_image[:, :, ::-1]
            height, width = original_image.shape[:2]
            image = self.aug.get_transform(original_image).apply_image(original_image)
            image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))

            inputs = []
            for _ in range(self.iterations):
                inputs.append({"image": image, "height": height, "width": width})

            batch_inputs = [inputs[i:i + self.batch_size] for i in range(0, self.iterations, self.batch_size)]  

            prediction_list = []
            for b in range(len(batch_inputs)):
                prediction_list.append(self.model(batch_inputs[b]))
            predictions = list(chain.from_iterable(prediction_list)) 

            return predictions


class MonteCarloDropoutHead:
    """
    Create a simple end-to-end predictor with the given config that runs on
    single device for a single input image.

    Compared to using the model directly, this class does the following additions:

    1. Load checkpoint from `cfg.MODEL.WEIGHTS`.
    2. Always take BGR image as the input and apply conversion defined by `cfg.INPUT.FORMAT`.
    3. Apply resizing defined by `cfg.INPUT.{MIN,MAX}_SIZE_TEST`.
    4. Take one input image and produce a single output, instead of a batch.

    If you'd like to do anything more fancy, please refer to its source code
    as examples to build and use the model manually.

    Attributes:
        metadata (Metadata): the metadata of the underlying dataset, obtained from
            cfg.DATASETS.TEST.

    Examples:
    ::
        pred = DefaultPredictor(cfg)
        inputs = cv2.imread("input.jpg")
        outputs = pred(inputs)
    """

    def __init__(self, cfg, iterations=1):
        self.cfg = cfg.clone()  # cfg can be modified by model
        self.model = build_model(self.cfg)
        self.model = patch_module(self.model)
        self.model.eval()
        if len(cfg.DATASETS.TEST):
            self.metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])

        checkpointer = DetectionCheckpointer(self.model)
        checkpointer.load(cfg.MODEL.WEIGHTS)

        self.aug = T.ResizeShortestEdge(
            [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST
        )

        self.input_format = cfg.INPUT.FORMAT
        assert self.input_format in ["RGB", "BGR"], self.input_format

        self.iterations = iterations

    def __call__(self, original_image):
        """
        Args:
            original_image (np.ndarray): an image of shape (H, W, C) (in BGR order).

        Returns:
            predictions (dict):
                the output of the model for one image only.
                See :doc:`/tutorials/models` for details about the format.
        """
        with torch.no_grad():  # https://github.com/sphinx-doc/sphinx/issues/4258
            # Apply pre-processing to image.
            if self.input_format == "RGB":
                # whether the model expects BGR inputs or RGB
                original_image = original_image[:, :, ::-1]
            height, width = original_image.shape[:2]
            image = self.aug.get_transform(original_image).apply_image(original_image)
            image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
            inputs = [{"image": image, "height": height, "width": width}]

            images = self.model.preprocess_image(inputs)
            features = self.model.backbone(images.tensor)
            proposals, _ = self.model.proposal_generator(images, features, None)

            features_ = [features[f] for f in self.model.roi_heads.box_in_features]
            box_features_pooler = self.model.roi_heads.box_pooler(features_, [x.proposal_boxes for x in proposals])
            
            prediction_list = []
            for _ in range(self.iterations):
                box_features = self.model.roi_heads.box_head(box_features_pooler)
                predictions = self.model.roi_heads.box_predictor(box_features)
                pred_instances, pred_inds = self.model.roi_heads.box_predictor.inference(predictions, proposals)
                pred_instances = self.model.roi_heads.forward_with_given_boxes(features, pred_instances)
                outputs = self.model._postprocess(pred_instances, inputs, images.image_sizes)
                prediction_list.append(outputs)
            predictions = list(chain.from_iterable(prediction_list)) 

            return predictions
# @Author: <NAME> # @Date:   2021-03-22 09:43:07 # @Last Modified by:   <NAME> # @Last Modified time: 2021-07-12 13:33:36 ## based on: detectron2.engine.defaults Create a simple end-to-end predictor with the given config that runs on     single device for a single input image.      Compared to using the model directly, this class does the following additions:      1. Load checkpoint from `cfg.MODEL.WEIGHTS`.     2. Always take BGR image as the input and apply conversion defined by `cfg.INPUT.FORMAT`.     3. Apply resizing defined by `cfg.INPUT.{MIN,MAX}_SIZE_TEST`.     4. Take one input image and produce a single output, instead of a batch.      If you'd like to do anything more fancy, please refer to its source code     as examples to build and use the model manually.      Attributes:         metadata (Metadata): the metadata of the underlying dataset, obtained from             cfg.DATASETS.TEST.      Examples:     ::         pred = DefaultPredictor(cfg)         inputs = cv2.imread("input.jpg")         outputs = pred(inputs) # cfg can be modified by model Args:             original_image (np.ndarray): an image of shape (H, W, C) (in BGR order).          Returns:             predictions (dict):                 the output of the model for one image only.                 See :doc:`/tutorials/models` for details about the format. # https://github.com/sphinx-doc/sphinx/issues/4258 # Apply pre-processing to image. # whether the model expects BGR inputs or RGB Create a simple end-to-end predictor with the given config that runs on     single device for a single input image.      Compared to using the model directly, this class does the following additions:      1. Load checkpoint from `cfg.MODEL.WEIGHTS`.     2. Always take BGR image as the input and apply conversion defined by `cfg.INPUT.FORMAT`.     3. Apply resizing defined by `cfg.INPUT.{MIN,MAX}_SIZE_TEST`.     4. Take one input image and produce a single output, instead of a batch.      If you'd like to do anything more fancy, please refer to its source code     as examples to build and use the model manually.      Attributes:         metadata (Metadata): the metadata of the underlying dataset, obtained from             cfg.DATASETS.TEST.      Examples:     ::         pred = DefaultPredictor(cfg)         inputs = cv2.imread("input.jpg")         outputs = pred(inputs) # cfg can be modified by model Args:             original_image (np.ndarray): an image of shape (H, W, C) (in BGR order).          Returns:             predictions (dict):                 the output of the model for one image only.                 See :doc:`/tutorials/models` for details about the format. # https://github.com/sphinx-doc/sphinx/issues/4258 # Apply pre-processing to image. # whether the model expects BGR inputs or RGB