text: # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
  import cv2
  import torch
  from torchvision import transforms as T
  from torchvision.transforms import functional as F
  from maskrcnn_benchmark.modeling.detector import build_detection_model
  from maskrcnn_benchmark.utils.checkpoint import DetectronCheckpointer
  from maskrcnn_benchmark.structures.image_list import to_image_list
  from maskrcnn_benchmark.modeling.roi_heads.mask_head.inference import Masker
  from maskrcnn_benchmark import layers as L
  from maskrcnn_benchmark.utils import cv2_util
  
  from image_decomposition import decompose_image
  from maskrcnn_benchmark.structures.boxlist_ops import *
  import math
  
  
  def assemble_patches(bboxes, patch_size, stride, image_size=(1080, 1920), type='xyxy'):
      # num_rows = math.ceil((image_size[0] - patch_size[0]) / stride[0]) + 1
      num_cols = math.ceil((image_size[1] - patch_size[1]) / stride[1]) + 1
      if type == 'xyxy':
          for patch_cnt, patch_bboxes in enumerate(bboxes):
              patch_bboxes.size = image_size
              mask = patch_bboxes.get_field("mask")
              if len(patch_bboxes) > 0:
                  top_left_x = (patch_cnt % num_cols) * stride[1]
                  top_left_y = (patch_cnt // num_cols) * stride[0]
                  patch_bboxes.bbox[:, 0] += top_left_x
                  patch_bboxes.bbox[:, 1] += top_left_y
                  patch_bboxes.bbox[:, 2] += top_left_x
                  patch_bboxes.bbox[:, 3] += top_left_y
                  im_mask = torch.zeros(image_size, dtype=torch.float)
                  y0 = max(0, top_left_y)
                  y1 = min(top_left_y+patch_size[0], image_size[0])
                  x0 = max(0, top_left_x)
                  x1 = min(top_left_x+patch_size[1], image_size[1])
                  im_mask[y0:y1, x0:x1] = mask[0][0][y0-top_left_y:y1-top_left_y, x0-top_left_x:x1-top_left_x]
                  patch_bboxes.extra_fields.pop("mask")
                  im_mask = torch.stack([im_mask], dim=0)[:, None]
                  patch_bboxes.add_field("mask", im_mask)
                  pass
              else:
                  mask = mask.new_empty((0, 1, image_size[0], image_size[1]))
                  patch_bboxes.add_field("mask", mask)
                  pass
  
      elif type == 'xywh':
          for patch_cnt, patch_bboxes in enumerate(bboxes):
              patch_bboxes.size = image_size
              if len(patch_bboxes) > 0:
                  top_left_x = (patch_cnt % num_cols) * stride[1]
                  top_left_y = (patch_cnt // num_cols) * stride[0]
                  patch_bboxes.bbox[:, 0] += top_left_x
                  patch_bboxes.bbox[:, 1] += top_left_y
  
      else:
          raise NotImplementedError(f'{type} not implemented.')
  
      return bboxes
  
  
  class Resize(object):
      def __init__(self, min_size, max_size):
          self.min_size = min_size
          self.max_size = max_size
  
      # modified from torchvision to add support for max size
      def get_size(self, image_size):
          w, h = image_size
          size = self.min_size
          max_size = self.max_size
          if max_size is not None:
              min_original_size = float(min((w, h)))
              max_original_size = float(max((w, h)))
              if max_original_size / min_original_size * size > max_size:
                  size = int(round(max_size * min_original_size / max_original_size))
  
          if (w <= h and w == size) or (h <= w and h == size):
              return (h, w)
  
          if w < h:
              ow = size
              oh = int(size * h / w)
          else:
              oh = size
              ow = int(size * w / h)
  
          return (oh, ow)
  
      def __call__(self, image):
          size = self.get_size(image.size)
          image = F.resize(image, size)
          return image
  
  
  class RIPDetector(object):
      # COCO categories for pretty print
      CATEGORIES = [
          "__background",
          "Flash Rip",
          "Rip Neck",
          "Rip Head",
          "Could be a Flash Rip?",
          "Could be a Rip?",
          "Rip"
      ]
  
      def __init__(
          self,
          cfg,
          confidence_threshold=0.7,
          show_mask_heatmaps=False,
          masks_per_dim=2,
          min_image_size=800,
          patch_size=(800, 800),
          patch_stride=(300, 700),
          weight_loading=None
      ):
          self.cfg = cfg.clone()
          self.model = build_detection_model(cfg)
          self.model.eval()
          self.device = torch.device(cfg.MODEL.DEVICE)
          self.model.to(self.device)
          self.min_image_size = min_image_size
          self.patch_size = patch_size
          self.patch_stride = patch_stride
          save_dir = cfg.OUTPUT_DIR
          checkpointer = DetectronCheckpointer(cfg, self.model, save_dir=save_dir)
          _ = checkpointer.load(cfg.MODEL.WEIGHT)
          
          if weight_loading:
              print('Loading weight from {}.'.format(weight_loading))
              _ = checkpointer._load_model(torch.load(weight_loading))
          
          self.transforms = self.build_transform()
  
          mask_threshold = -1 if show_mask_heatmaps else 0.5
          self.masker = Masker(threshold=mask_threshold, padding=1)
  
          # used to make colors for each class
          self.palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
  
          self.cpu_device = torch.device("cpu")
          self.confidence_threshold = confidence_threshold
          self.show_mask_heatmaps = show_mask_heatmaps
          self.masks_per_dim = masks_per_dim
  
      def build_transform(self):
          """
          Creates a basic transformation that was used to train the models
          """
          cfg = self.cfg
  
          # we are loading images with OpenCV, so we don't need to convert them
          # to BGR, they are already! So all we need to do is to normalize
          # by 255 if we want to convert to BGR255 format, or flip the channels
          # if we want it to be in RGB in [0-1] range.
          if cfg.INPUT.TO_BGR255:
              to_bgr_transform = T.Lambda(lambda x: x * 255)
          else:
              to_bgr_transform = T.Lambda(lambda x: x[[2, 1, 0]])
  
          normalize_transform = T.Normalize(
              mean=cfg.INPUT.PIXEL_MEAN, std=cfg.INPUT.PIXEL_STD
          )
          min_size = cfg.INPUT.MIN_SIZE_TEST
          max_size = cfg.INPUT.MAX_SIZE_TEST
          transform = T.Compose(
              [
                  T.ToPILImage(),
                  Resize(min_size, max_size),
                  T.ToTensor(),
                  to_bgr_transform,
                  normalize_transform,
              ]
          )
          return transform
  
      def run_on_opencv_image(self, image, thickness=1):
          """
          Arguments:
              image (np.ndarray): an image as returned by OpenCV
  
          Returns:
              prediction (BoxList): the detected objects. Additional information
                  of the detection properties can be found in the fields of
                  the BoxList via `prediction.fields()`
          """
          predictions = self.compute_prediction(image)
          top_predictions = self.select_top_predictions(predictions)
  
          result = image.copy()
          if self.show_mask_heatmaps:
              return self.create_mask_montage(result, top_predictions)
          result = self.overlay_boxes(result, top_predictions, thickness)
          if self.cfg.MODEL.MASK_ON:
              result = self.overlay_mask(result, top_predictions)
          result = self.overlay_class_names(result, top_predictions, thickness)
  
          return result
  
      def run_on_image_patches(self, original_image, thickness=1):
          patches = decompose_image(original_image, None, self.patch_size, self.patch_stride)
          detection_list = []
          for k, patch in patches.items():
              predictions = self.compute_prediction(patch.image)
              top_predictions = self.select_top_predictions(predictions)
              detection_list.append(top_predictions)
          detection_list = assemble_patches(detection_list, self.patch_size, self.patch_stride)
          detections = cat_boxlist(detection_list)
          detections = boxlist_nms(detections, nms_thresh=0.1)
          detections = self.select_top_predictions(detections)
          detections.add_field("mask", detections.get_field("mask").to(torch.uint8))
  
          result = original_image.copy()
          if self.show_mask_heatmaps:
              return self.create_mask_montage(result, detections)
          result = self.overlay_boxes(result, detections, thickness)
          if self.cfg.MODEL.MASK_ON:
              result = self.overlay_mask(result, detections)
          result = self.overlay_class_names(result, detections, thickness)
  
          return result
  
      def compute_prediction(self, original_image):
          """
          Arguments:
              original_image (np.ndarray): an image as returned by OpenCV
  
          Returns:
              prediction (BoxList): the detected objects. Additional information
                  of the detection properties can be found in the fields of
                  the BoxList via `prediction.fields()`
          """
          # apply pre-processing to image
          image = self.transforms(original_image)
          # convert to an ImageList, padded so that it is divisible by
          # cfg.DATALOADER.SIZE_DIVISIBILITY
          image_list = to_image_list(image, self.cfg.DATALOADER.SIZE_DIVISIBILITY)
          image_list = image_list.to(self.device)
          # compute predictions
          with torch.no_grad():
              predictions = self.model(image_list)
          predictions = [o.to(self.cpu_device) for o in predictions]
  
          # always single image is passed at a time
          prediction = predictions[0]
  
          # reshape prediction (a BoxList) into the original image size
          height, width = original_image.shape[:-1]
          prediction = prediction.resize((width, height))
  
          if prediction.has_field("mask"):
              # if we have masks, paste the masks in the right position
              # in the image, as defined by the bounding boxes
              masks = prediction.get_field("mask")
              # always single image is passed at a time
              masks = self.masker([masks], [prediction])[0]
              prediction.add_field("mask", masks)
          return prediction
  
      def select_top_predictions(self, predictions):
          """
          Select only predictions which have a `score` > self.confidence_threshold,
          and returns the predictions in descending order of score
  
          Arguments:
              predictions (BoxList): the result of the computation by the model.
                  It should contain the field `scores`.
  
          Returns:
              prediction (BoxList): the detected objects. Additional information
                  of the detection properties can be found in the fields of
                  the BoxList via `prediction.fields()`
          """
          scores = predictions.get_field("scores")
          keep = torch.nonzero(scores > self.confidence_threshold).squeeze(1)
          predictions = predictions[keep]
          scores = predictions.get_field("scores")
          _, idx = scores.sort(0, descending=True)
          return predictions[idx]
  
      def compute_colors_for_labels(self, labels):
          """
          Simple function that adds fixed colors depending on the class
          """
          colors = labels[:, None] * self.palette
          colors = (colors % 255).numpy().astype("uint8")
          return colors
  
      def overlay_boxes(self, image, predictions, thickness=1):
          """
          Adds the predicted boxes on top of the image
  
          Arguments:
              image (np.ndarray): an image as returned by OpenCV
              predictions (BoxList): the result of the computation by the model.
                  It should contain the field `labels`.
          """
          labels = predictions.get_field("labels")
          boxes = predictions.bbox
  
          colors = self.compute_colors_for_labels(labels).tolist()
  
          for box, color in zip(boxes, colors):
              box = box.to(torch.int64)
              top_left, bottom_right = box[:2].tolist(), box[2:].tolist()
              image = cv2.rectangle(
                  image, tuple(top_left), tuple(bottom_right), tuple(color), thickness
              )
  
          return image
  
      def overlay_mask(self, image, predictions):
          """
          Adds the instances contours for each predicted object.
          Each label has a different color.
  
          Arguments:
              image (np.ndarray): an image as returned by OpenCV
              predictions (BoxList): the result of the computation by the model.
                  It should contain the field `mask` and `labels`.
          """
          masks = predictions.get_field("mask").numpy()
          labels = predictions.get_field("labels")
  
          colors = self.compute_colors_for_labels(labels).tolist()
  
          for mask, color in zip(masks, colors):
              thresh = mask[0, :, :, None]
              contours, hierarchy = cv2_util.findContours(
                  thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
              )
              image = cv2.drawContours(image, contours, -1, color, 3)
  
          composite = image
  
          return composite
  
      def create_mask_montage(self, image, predictions):
          """
          Create a montage showing the probability heatmaps for each one one of the
          detected objects
  
          Arguments:
              image (np.ndarray): an image as returned by OpenCV
              predictions (BoxList): the result of the computation by the model.
                  It should contain the field `mask`.
          """
          masks = predictions.get_field("mask")
          masks_per_dim = self.masks_per_dim
          masks = L.interpolate(
              masks.float(), scale_factor=1 / masks_per_dim
          ).byte()
          height, width = masks.shape[-2:]
          max_masks = masks_per_dim ** 2
          masks = masks[:max_masks]
          # handle case where we have less detections than max_masks
          if len(masks) < max_masks:
              masks_padded = torch.zeros(max_masks, 1, height, width, dtype=torch.uint8)
              masks_padded[: len(masks)] = masks
              masks = masks_padded
          masks = masks.reshape(masks_per_dim, masks_per_dim, height, width)
          result = torch.zeros(
              (masks_per_dim * height, masks_per_dim * width), dtype=torch.uint8
          )
          for y in range(masks_per_dim):
              start_y = y * height
              end_y = (y + 1) * height
              for x in range(masks_per_dim):
                  start_x = x * width
                  end_x = (x + 1) * width
                  result[start_y:end_y, start_x:end_x] = masks[y, x]
          return cv2.applyColorMap(result.numpy(), cv2.COLORMAP_JET)
  
      def overlay_class_names(self, image, predictions, thickness=1):
          """
          Adds detected class names and scores in the positions defined by the
          top-left corner of the predicted bounding box
  
          Arguments:
              image (np.ndarray): an image as returned by OpenCV
              predictions (BoxList): the result of the computation by the model.
                  It should contain the field `scores` and `labels`.
          """
          scores = predictions.get_field("scores").tolist()
          labels = predictions.get_field("labels").tolist()
          labels = [self.CATEGORIES[i] for i in labels]
          boxes = predictions.bbox
  
          template = "{}: {:.2f}"
          for box, score, label in zip(boxes, scores, labels):
              x, y = box[:2]
              s = template.format(label, score)
              cv2.putText(
                  image, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, .5*thickness, (255, 255, 255), thickness
              )
  
          return image
  