import torch
import detectron2.data.transforms as T

from detectron2.structures import Instances
from detectron2.layers import paste_masks_in_image
from detectron2.modeling import build_model
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.data import MetadataCatalog
from detectron2.checkpoint import DetectionCheckpointer

def detector_postprocess(results, output_height, output_width, mask_threshold=0.5):
    """
    In addition to the post processing of detectron2, we add scalign for 
    bezier control points.
    """
    scale_x, scale_y = (output_width / results.image_size[1], output_height / results.image_size[0])
    # Change to 'if is_tracing' after PT1.7
    if isinstance(output_height, torch.Tensor):
        # Converts integer tensors to float temporaries to ensure true
        # division is performed when computing scale_x and scale_y.
        output_width_tmp = output_width.float()
        output_height_tmp = output_height.float()
        new_size = torch.stack([output_height, output_width])
    else:
        new_size = (output_height, output_width)
        output_width_tmp = output_width
        output_height_tmp = output_height

    scale_x, scale_y = (
        output_width_tmp / results.image_size[1],
        output_height_tmp / results.image_size[0],
    )
    results = Instances(new_size, **results.get_fields())

    if results.has("pred_boxes"):
        output_boxes = results.pred_boxes
    elif results.has("proposal_boxes"):
        output_boxes = results.proposal_boxes
    else:
        output_boxes = None
    assert output_boxes is not None, "Predictions must contain boxes!"

    output_boxes.scale(scale_x, scale_y)
    output_boxes.clip(results.image_size)

    results = results[output_boxes.nonempty()]

    if results.has("pred_visible_masks"):
        results.pred_visible_masks = paste_masks_in_image(
            results.pred_visible_masks[:, 0, :, :],
            results.pred_boxes,
            results.image_size,
            threshold=mask_threshold,
        )
        if results.pred_visible_masks.shape[0] == 0:
            results.pred_visible_masks = torch.zeros_like(results.pred_masks)

    if results.has("pred_occluded_masks"):
        results.pred_occluded_masks = paste_masks_in_image(
            results.pred_occluded_masks[:, 0, :, :],
            results.pred_boxes,
            results.image_size,
            threshold=mask_threshold,
        )
    elif results.has("pred_visible_masks"):
        results.pred_occluded_masks = torch.bitwise_xor(results.pred_masks, results.pred_visible_masks)
        if results.has("pred_occlusions"):
            for i, pred_occlusion in enumerate(results.pred_occlusions):
                if int(pred_occlusion) == 0:
                    results.pred_occluded_masks[i] = torch.zeros_like(results.pred_occluded_masks[i])
        else:
            occlusions = []
            for i in range(results.pred_masks.shape[0]):
                ratio = torch.sum(results.pred_visible_masks[i]) / torch.sum(results.pred_masks[i])
                if ratio > 0.95:
                    occlusions.append(0)
                else:
                    occlusions.append(1)
            results.pred_occlusions = torch.Tensor(occlusions)
        
    # scale bezier points
    if results.has("beziers"):
        beziers = results.beziers
        # scale and clip in place
        beziers[:, 0::2] *= scale_x
        beziers[:, 1::2] *= scale_y
        h, w = results.image_size
        beziers[:, 0].clamp_(min=0, max=w)
        beziers[:, 1].clamp_(min=0, max=h)
        beziers[:, 6].clamp_(min=0, max=w)
        beziers[:, 7].clamp_(min=0, max=h)
        beziers[:, 8].clamp_(min=0, max=w)
        beziers[:, 9].clamp_(min=0, max=h)
        beziers[:, 14].clamp_(min=0, max=w)
        beziers[:, 15].clamp_(min=0, max=h)

    return results


class DefaultPredictor:
    """
    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.
    This is meant for simple demo purposes, so it does the above steps automatically.
    This is not meant for benchmarks or running complicated inference logic.
    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):
        self.cfg = cfg.clone()  # cfg can be modified by model
        self.model = build_model(self.cfg)
        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.input_format = cfg.INPUT.FORMAT
        assert self.input_format in ["RGB", "BGR"], self.input_format

    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 = original_image
            # 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}
            predictions = self.model([inputs])[0]
            return predictions
import torch
import detectron2.data.transforms as T

from detectron2.structures import Instances
from detectron2.layers import paste_masks_in_image
from detectron2.modeling import build_model
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.data import MetadataCatalog
from detectron2.checkpoint import DetectionCheckpointer

def detector_postprocess(results, output_height, output_width, mask_threshold=0.5):
    """
    In addition to the post processing of detectron2, we add scalign for 
    bezier control points.
    """
    scale_x, scale_y = (output_width / results.image_size[1], output_height / results.image_size[0])
    # Change to 'if is_tracing' after PT1.7
    if isinstance(output_height, torch.Tensor):
        # Converts integer tensors to float temporaries to ensure true
        # division is performed when computing scale_x and scale_y.
        output_width_tmp = output_width.float()
        output_height_tmp = output_height.float()
        new_size = torch.stack([output_height, output_width])
    else:
        new_size = (output_height, output_width)
        output_width_tmp = output_width
        output_height_tmp = output_height

    scale_x, scale_y = (
        output_width_tmp / results.image_size[1],
        output_height_tmp / results.image_size[0],
    )
    results = Instances(new_size, **results.get_fields())

    if results.has("pred_boxes"):
        output_boxes = results.pred_boxes
    elif results.has("proposal_boxes"):
        output_boxes = results.proposal_boxes
    else:
        output_boxes = None
    assert output_boxes is not None, "Predictions must contain boxes!"

    output_boxes.scale(scale_x, scale_y)
    output_boxes.clip(results.image_size)

    results = results[output_boxes.nonempty()]

    if results.has("pred_visible_masks"):
        results.pred_visible_masks = paste_masks_in_image(
            results.pred_visible_masks[:, 0, :, :],
            results.pred_boxes,
            results.image_size,
            threshold=mask_threshold,
        )
        if results.pred_visible_masks.shape[0] == 0:
            results.pred_visible_masks = torch.zeros_like(results.pred_masks)

    if results.has("pred_occluded_masks"):
        results.pred_occluded_masks = paste_masks_in_image(
            results.pred_occluded_masks[:, 0, :, :],
            results.pred_boxes,
            results.image_size,
            threshold=mask_threshold,
        )
    elif results.has("pred_visible_masks"):
        results.pred_occluded_masks = torch.bitwise_xor(results.pred_masks, results.pred_visible_masks)
        if results.has("pred_occlusions"):
            for i, pred_occlusion in enumerate(results.pred_occlusions):
                if int(pred_occlusion) == 0:
                    results.pred_occluded_masks[i] = torch.zeros_like(results.pred_occluded_masks[i])
        else:
            occlusions = []
            for i in range(results.pred_masks.shape[0]):
                ratio = torch.sum(results.pred_visible_masks[i]) / torch.sum(results.pred_masks[i])
                if ratio > 0.95:
                    occlusions.append(0)
                else:
                    occlusions.append(1)
            results.pred_occlusions = torch.Tensor(occlusions)
        
    # scale bezier points
    if results.has("beziers"):
        beziers = results.beziers
        # scale and clip in place
        beziers[:, 0::2] *= scale_x
        beziers[:, 1::2] *= scale_y
        h, w = results.image_size
        beziers[:, 0].clamp_(min=0, max=w)
        beziers[:, 1].clamp_(min=0, max=h)
        beziers[:, 6].clamp_(min=0, max=w)
        beziers[:, 7].clamp_(min=0, max=h)
        beziers[:, 8].clamp_(min=0, max=w)
        beziers[:, 9].clamp_(min=0, max=h)
        beziers[:, 14].clamp_(min=0, max=w)
        beziers[:, 15].clamp_(min=0, max=h)

    return results


class DefaultPredictor:
    """
    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.
    This is meant for simple demo purposes, so it does the above steps automatically.
    This is not meant for benchmarks or running complicated inference logic.
    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):
        self.cfg = cfg.clone()  # cfg can be modified by model
        self.model = build_model(self.cfg)
        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.input_format = cfg.INPUT.FORMAT
        assert self.input_format in ["RGB", "BGR"], self.input_format

    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 = original_image
            # 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}
            predictions = self.model([inputs])[0]
            return predictions
In addition to the post processing of detectron2, we add scalign for      bezier control points. # Change to 'if is_tracing' after PT1.7 # Converts integer tensors to float temporaries to ensure true # division is performed when computing scale_x and scale_y. # scale bezier points # scale and clip in place 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.     This is meant for simple demo purposes, so it does the above steps automatically.     This is not meant for benchmarks or running complicated inference logic.     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 # image = self.aug.get_transform(original_image).apply_image(original_image)