text: # -*- coding: utf-8 -*-
  # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
  
  import torch
  from torch import nn
  
  from detectron2.structures import ImageList
  
  from ..backbone import build_backbone
  from ..postprocessing import detector_postprocess, sem_seg_postprocess, detector_resize_logits
  from ..proposal_generator import build_proposal_generator
  from ..roi_heads import build_roi_heads
  from .build import META_ARCH_REGISTRY
  from .semantic_seg import build_sem_seg_head
  
  from pytorch_permuto.pytorch_bcrf import PyTorchBCRF
  from crf.densecrf import DenseCRFParams
  
  from ..roi_heads.mask_head import mask_rcnn_loss2
  import copy
  
  __all__ = ["PanopticFPN"]
  
  
  @META_ARCH_REGISTRY.register()
  class PanopticFPN(nn.Module):
      """
      Main class for Panoptic FPN architectures (see https://arxiv.org/abs/1901.02446).
      """
  
      def __init__(self, cfg):
          super().__init__()
  
          self.device = torch.device(cfg.MODEL.DEVICE)
  
          self.instance_loss_weight = cfg.MODEL.PANOPTIC_FPN.INSTANCE_LOSS_WEIGHT
  
          # options when combining instance & semantic outputs
          self.combine_on = cfg.MODEL.PANOPTIC_FPN.COMBINE.ENABLED
          self.combine_overlap_threshold = cfg.MODEL.PANOPTIC_FPN.COMBINE.OVERLAP_THRESH
          self.combine_stuff_area_limit = cfg.MODEL.PANOPTIC_FPN.COMBINE.STUFF_AREA_LIMIT
          self.combine_instances_confidence_threshold = (
              cfg.MODEL.PANOPTIC_FPN.COMBINE.INSTANCES_CONFIDENCE_THRESH
          )
  
          self.backbone = build_backbone(cfg)
          self.proposal_generator = build_proposal_generator(cfg, self.backbone.output_shape())
          self.roi_heads = build_roi_heads(cfg, self.backbone.output_shape())
          self.sem_seg_head = build_sem_seg_head(cfg, self.backbone.output_shape())
  
          pixel_mean = torch.Tensor(cfg.MODEL.PIXEL_MEAN).to(self.device).view(3, 1, 1)
          pixel_std = torch.Tensor(cfg.MODEL.PIXEL_STD).to(self.device).view(3, 1, 1)
          self.normalizer = lambda x: (x - pixel_mean) / pixel_std
  
          sem_params = DenseCRFParams(spatial_ker_weight=0.5, bilateral_ker_weight=0.5, alpha=100, beta=1, gamma=3)
          ins_params = DenseCRFParams(spatial_ker_weight=0.5, bilateral_ker_weight=0.5, alpha=100, beta=2, gamma=3)
  
          bcrf_num_labels = cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES+cfg.MODEL.ROI_HEADS.NUM_CLASSES-1
          bcrf_thing_labels = tuple(range(cfg.MODEL.ROI_HEADS.NUM_CLASSES))
          bcrf_stuff_labels = tuple(range(cfg.MODEL.ROI_HEADS.NUM_CLASSES,bcrf_num_labels))
          self.bcrf = PyTorchBCRF(sem_params=sem_params, ins_params=ins_params, num_labels=bcrf_num_labels,
                                  thing_labels=bcrf_thing_labels,
                                  stuff_labels=bcrf_stuff_labels,
                                  num_iterations=5)
          self.bcrf_bg_label = cfg.MODEL.ROI_HEADS.NUM_CLASSES
          self.bcrf_down_factor = cfg.MODEL.BCRF_HEAD.DOWN_FACTOR
          self.to(self.device)
  
      def forward(self, batched_inputs):
          """
          Args:
              batched_inputs: a list, batched outputs of :class:`DatasetMapper`.
                  Each item in the list contains the inputs for one image.
  
                  For now, each item in the list is a dict that contains:
  
                  * "image": Tensor, image in (C, H, W) format.
                  * "instances": Instances
                  * "sem_seg": semantic segmentation ground truth.
                  * Other information that's included in the original dicts, such as:
                    "height", "width" (int): the output resolution of the model, used in inference.
                    See :meth:`postprocess` for details.
  
          Returns:
              list[dict]:
                  each dict is the results for one image. The dict contains the following keys:
  
                  * "instances": see :meth:`GeneralizedRCNN.forward` for its format.
                  * "sem_seg": see :meth:`SemanticSegmentor.forward` for its format.
                  * "panoptic_seg": available when `PANOPTIC_FPN.COMBINE.ENABLED`.
                    See the return value of
                    :func:`combine_semantic_and_instance_outputs` for its format.
          """
          #print(batched_inputs[0]['file_name'])
          images = [x["image"].to(self.device) for x in batched_inputs]
          images = [self.normalizer(x) for x in images]
          images = ImageList.from_tensors(images, self.backbone.size_divisibility)
          features = self.backbone(images.tensor)
  
          if "proposals" in batched_inputs[0]:
              proposals = [x["proposals"].to(self.device) for x in batched_inputs]
              proposal_losses = {}
  
          if "sem_seg" in batched_inputs[0]:
              gt_sem_seg = [x["sem_seg"].to(self.device) for x in batched_inputs]
              gt_sem_seg = ImageList.from_tensors(
                  gt_sem_seg, self.backbone.size_divisibility, self.sem_seg_head.ignore_value
              ).tensor
          else:
              gt_sem_seg = None
          sem_seg_results, sem_seg_losses = self.sem_seg_head(features, gt_sem_seg)
  
          if "instances" in batched_inputs[0]:
              gt_instances = [x["instances"].to(self.device) for x in batched_inputs]
          else:
              gt_instances = None
          if self.proposal_generator:
              proposals, proposal_losses = self.proposal_generator(images, features, gt_instances)
          detector_results, detector_losses, ins_logits, gt_masks, gt_classes, mask_loss = self.roi_heads(
              images, features, proposals, gt_instances
          )
          print(ins_logits.shape, sem_seg_results.shape)
          pt_sem_seg_results = []
          pt_detector_results = []
          for sem_seg_result, detector_result, input_per_image, image_size in zip(
                  sem_seg_results, detector_results, batched_inputs, images.image_sizes
          ):
              height = input_per_image.get("height", image_size[0])//self.bcrf_down_factor
              width = input_per_image.get("width", image_size[1])//self.bcrf_down_factor
              det_template = copy.deepcopy(detector_result)
  
              sem_seg_r = sem_seg_postprocess(sem_seg_result, image_size, height, width)
              ins_logits = detector_resize_logits(detector_result, ins_logits.sigmoid(), height, width)
              
              image_r = sem_seg_postprocess(input_per_image.get("image").to(dtype=torch.float32), image_size, height,
                                            width).to(torch.device("cpu"))
  
              # .................................................................................
              bcrf_device = image_r.device
              sem_seg_r = sem_seg_r.to(bcrf_device)
              ins_logits = ins_logits.to(bcrf_device)
              if self.training:
                  obj_classes = gt_classes
              else:
                  obj_classes = det_template._fields.get('pred_classes')
              obj_classes = obj_classes.to(bcrf_device)
              obj_classes = torch.cat((obj_classes, torch.Tensor([self.bcrf_bg_label]).type(torch.int64)), dim=0).to(
                  sem_seg_result.device)
  
              if ins_logits.shape[0] == 0:
                  back_ground = torch.ones(ins_logits.shape[1:]) * 11  # max 98 % confidence
              else:
                  ins_logits = -torch.log(-1 + 1/ins_logits)
                  temp_ins_logits = ins_logits.permute((1, 2, 0))
                  back_ground, _ = torch.max(temp_ins_logits, dim=2)
                  back_ground = - back_ground
              back_ground = torch.unsqueeze(back_ground, dim=0)
              back_ground = back_ground.to(bcrf_device)
              ins_logits = torch.cat((ins_logits.to(dtype=back_ground.dtype), back_ground), dim=0)
              sem_seg_r, ins_logits = self.bcrf(image_r, sem_seg_r, ins_logits, obj_classes)
              ins_logits = ins_logits[:-1]
              # ..................................................................................
              if self.training:
                  losses = {}
                  sem_seg_r = sem_seg_postprocess(sem_seg_r, sem_seg_r.shape[1:], sem_seg_result.shape[1],
                                                  sem_seg_result.shape[2])
                  sem_seg_losses = self.sem_seg_head.losses(torch.unsqueeze(sem_seg_r, 0).to(gt_sem_seg.device),
                                                            gt_sem_seg)
                  if gt_masks is not None:
                      new_mask_loss = mask_rcnn_loss2(ins_logits, gt_masks, detector_result.proposal_boxes)
                  else:
                      new_mask_loss = {"loss_mask": 0}
                  detector_losses.update(new_mask_loss)
                  losses.update(sem_seg_losses)
                  losses.update({k: v * self.instance_loss_weight for k, v in detector_losses.items()})
                  losses.update(proposal_losses)
                  return losses
              else:
                  height = input_per_image.get("height", image_size[0])
                  width = input_per_image.get("width", image_size[1])
                  detector_r = detector_postprocess(det_template, height, width)
                  if ins_logits.shape[0] != 0:
                      ins_logits = sem_seg_postprocess(ins_logits, ins_logits.shape[1:],height,
                                                      width)
                      ins_probs = ins_logits.sigmoid()
                      detector_r.pred_masks = ins_probs >= 0.5
                  sem_seg_r = sem_seg_postprocess(sem_seg_r, sem_seg_r.shape[1:], height,
                                                  width)
                  #ins_probs = ins_logits.sigmoid()
                  #detector_r.pred_masks = ins_probs >= 0.5
                  pt_sem_seg_results.append(sem_seg_r)
                  pt_detector_results.append(detector_r)
  
          processed_results = []
          for sem_seg_result, detector_result, input_per_image, image_size in zip(
                  pt_sem_seg_results, pt_detector_results, batched_inputs, images.image_sizes
          ):
              # height = input_per_image.get("height", image_size[0])
              # width = input_per_image.get("width", image_size[1])
              # sem_seg_r = sem_seg_postprocess(sem_seg_result, image_size, height, width)
              # detector_r = detector_postprocess(detector_result, height, width)
  
              processed_results.append({"sem_seg": sem_seg_r, "instances": detector_r})
  
              if self.combine_on:
                  panoptic_r = combine_semantic_and_instance_outputs(
                      detector_r,
                      sem_seg_r.argmax(dim=0),
                      self.combine_overlap_threshold,
                      self.combine_stuff_area_limit,
                      self.combine_instances_confidence_threshold,
                  )
                  processed_results[-1]["panoptic_seg"] = panoptic_r
          return processed_results
  
  
  def combine_semantic_and_instance_outputs(
          instance_results,
          semantic_results,
          overlap_threshold,
          stuff_area_limit,
          instances_confidence_threshold,
  ):
      """
      Implement a simple combining logic following
      "combine_semantic_and_instance_predictions.py" in panopticapi
      to produce panoptic segmentation outputs.
  
      Args:
          instance_results: output of :func:`detector_postprocess`.
          semantic_results: an (H, W) tensor, each is the contiguous semantic
              category id
  
      Returns:
          panoptic_seg (Tensor): of shape (height, width) where the values are ids for each segment.
          segments_info (list[dict]): Describe each segment in `panoptic_seg`.
              Each dict contains keys "id", "category_id", "isthing".
      """
      panoptic_seg = torch.zeros_like(semantic_results, dtype=torch.int32)
  
      # sort instance outputs by scores
      sorted_inds = torch.argsort(-instance_results.scores)
  
      current_segment_id = 0
      segments_info = []
  
      instance_masks = instance_results.pred_masks.to(dtype=torch.bool, device=panoptic_seg.device)
  
      # Add instances one-by-one, check for overlaps with existing ones
      for inst_id in sorted_inds:
          score = instance_results.scores[inst_id].item()
          if score < instances_confidence_threshold:
              break
          mask = instance_masks[inst_id]  # H,W
          mask_area = mask.sum().item()
  
          if mask_area == 0:
              continue
  
          intersect = (mask > 0) & (panoptic_seg > 0)
          intersect_area = intersect.sum().item()
  
          if intersect_area * 1.0 / mask_area > overlap_threshold:
              continue
  
          if intersect_area > 0:
              mask = mask & (panoptic_seg == 0)
  
          current_segment_id += 1
          panoptic_seg[mask] = current_segment_id
          segments_info.append(
              {
                  "id": current_segment_id,
                  "isthing": True,
                  "score": score,
                  "category_id": instance_results.pred_classes[inst_id].item(),
                  "instance_id": inst_id.item(),
              }
          )
  
      # Add semantic results to remaining empty areas
      semantic_labels = torch.unique(semantic_results).cpu().tolist()
      for semantic_label in semantic_labels:
          if semantic_label == 0:  # 0 is a special "thing" class
              continue
          mask = (semantic_results == semantic_label) & (panoptic_seg == 0)
          mask_area = mask.sum().item()
          if mask_area < stuff_area_limit:
              continue
  
          current_segment_id += 1
          panoptic_seg[mask] = current_segment_id
          segments_info.append(
              {
                  "id": current_segment_id,
                  "isthing": False,
                  "category_id": semantic_label,
                  "area": mask_area,
              }
          )
  
      return panoptic_seg, segments_info
  