text: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  from mmseg.core import add_prefix
  from mmseg.ops import resize
  from .. import builder
  from ..builder import SEGMENTORS
  from .base import BaseSegmentor
  from ..cluster.kmeans import kmeans_debackground
  import numpy as np
  @SEGMENTORS.register_module()
  class EncoderDecoder(BaseSegmentor):
      """Encoder Decoder segmentors.
  
      EncoderDecoder typically consists of backbone, decode_head, auxiliary_head.
      Note that auxiliary_head is only used for deep supervision during training,
      which could be dumped during inference.
      """
  
      def __init__(self,
                   backbone,
                   decode_head,
                   neck=None,
                   auxiliary_head=None,
                   train_cfg=None,
                   test_cfg=None,
                   pretrained=None):
          super(EncoderDecoder, self).__init__()
          self.backbone = builder.build_backbone(backbone)
          if neck is not None:
              self.neck = builder.build_neck(neck)
          self._init_decode_head(decode_head)
    #     self._init_auxiliary_head(auxiliary_head)
  
          self.train_cfg = train_cfg
          self.test_cfg = test_cfg
  
          self.init_weights(pretrained=pretrained)
  
          assert self.with_decode_head
  
      def _init_decode_head(self, decode_head):
          """Initialize ``decode_head``"""
          self.decode_head = builder.build_head(decode_head)
          self.align_corners = self.decode_head.align_corners
          self.num_classes = self.decode_head.num_classes
  
      def _init_auxiliary_head(self, auxiliary_head):
          """Initialize ``auxiliary_head``"""
          if auxiliary_head is not None:
              if isinstance(auxiliary_head, list):
                  self.auxiliary_head = nn.ModuleList()
                  for head_cfg in auxiliary_head:
                      self.auxiliary_head.append(builder.build_head(head_cfg))
              else:
                  self.auxiliary_head = builder.build_head(auxiliary_head)
  
      def init_weights(self, pretrained=None):
          """Initialize the weights in backbone and heads.
  
          Args:
              pretrained (str, optional): Path to pre-trained weights.
                  Defaults to None.
          """
  
          super(EncoderDecoder, self).init_weights(pretrained)
          self.backbone.init_weights(pretrained=pretrained)
          self.decode_head.init_weights()
          if self.with_auxiliary_head:
              if isinstance(self.auxiliary_head, nn.ModuleList):
                  for aux_head in self.auxiliary_head:
                      aux_head.init_weights()
              else:
                  self.auxiliary_head.init_weights()
  
      def extract_feat(self, img):
          """Extract features from images."""
          x = self.backbone(img)
          if self.with_neck:
              x = self.neck(x)
          return x
  
      def encode_decode(self, img, img_metas):
          """Encode images with backbone and decode into a semantic segmentation
          map of the same size as input."""
          x = self.extract_feat(img)
          # print('x',x)
          # print(np.size(x))
  
          out = self._decode_head_forward_test(x, img_metas)
          # print('out',out)
          # print(np.size(out))
          out = resize(
              input=out,
              size=img.shape[2:],
              mode='bilinear',
              align_corners=self.align_corners)
          # print('fout',out)
          # print(np.size(out))
          return out
  
      def _decode_head_forward_train(self, x, img_metas, gt_semantic_seg):
          """Run forward function and calculate loss for decode head in
          training."""
          losses = dict()
          loss_decode = self.decode_head.forward_train(x, img_metas,
                                                       gt_semantic_seg,
                                                       self.train_cfg)
  
          losses.update(add_prefix(loss_decode, 'decode'))
          return losses
  
      def _decode_head_forward_test(self, x, img_metas):
          """Run forward function and calculate loss for decode head in
          inference."""
          seg_logits = self.decode_head.forward_test(x, img_metas, self.test_cfg)
          return seg_logits
  
      def _auxiliary_head_forward_train(self, x, img_metas, gt_semantic_seg):
          """Run forward function and calculate loss for auxiliary head in
          training."""
          losses = dict()
          if isinstance(self.auxiliary_head, nn.ModuleList):
              for idx, aux_head in enumerate(self.auxiliary_head):
                  loss_aux = aux_head.forward_train(x, img_metas,
                                                    gt_semantic_seg,
                                                    self.train_cfg)
                  losses.update(add_prefix(loss_aux, f'aux_{idx}'))
          else:
              loss_aux = self.auxiliary_head.forward_train(
                  x, img_metas, gt_semantic_seg, self.train_cfg)
              losses.update(add_prefix(loss_aux, 'aux'))
  
          return losses
  
      def forward_dummy(self, img):
          """Dummy forward function."""
          seg_logit = self.encode_decode(img, None)
  
          return seg_logit
  
      def forward_train(self, img, img_metas, gt_semantic_seg):
          """Forward function for training.
  
          Args:
              img (Tensor): Input images.
              img_metas (list[dict]): List of image info dict where each dict
                  has: 'img_shape', 'scale_factor', 'flip', and may also contain
                  'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
                  For details on the values of these keys see
                  `mmseg/datasets/pipelines/formatting.py:Collect`.
              gt_semantic_seg (Tensor): Semantic segmentation masks
                  used if the architecture supports semantic segmentation task.
  
          Returns:
              dict[str, Tensor]: a dictionary of loss components
          """
  
          x = self.extract_feat(img)
  
          losses = dict()
  
          loss_decode = self._decode_head_forward_train(x, img_metas,
                                                        gt_semantic_seg)
          losses.update(loss_decode)
  
          if self.with_auxiliary_head:
              loss_aux = self._auxiliary_head_forward_train(
                  x, img_metas, gt_semantic_seg)
              losses.update(loss_aux)
  
          return losses
  
      # TODO refactor
      def slide_inference(self, img, img_meta, rescale):
          """Inference by sliding-window with overlap.
  
          If h_crop > h_img or w_crop > w_img, the small patch will be used to
          decode without padding.
          """
  
          h_stride, w_stride = self.test_cfg.stride
          h_crop, w_crop = self.test_cfg.crop_size
          batch_size, _, h_img, w_img = img.size()
          num_classes = self.num_classes
          h_grids = max(h_img - h_crop + h_stride - 1, 0) // h_stride + 1
          w_grids = max(w_img - w_crop + w_stride - 1, 0) // w_stride + 1
          preds = img.new_zeros((batch_size, num_classes, h_img, w_img))
          count_mat = img.new_zeros((batch_size, 1, h_img, w_img))
          for h_idx in range(h_grids):
              for w_idx in range(w_grids):
                  y1 = h_idx * h_stride
                  x1 = w_idx * w_stride
                  y2 = min(y1 + h_crop, h_img)
                  x2 = min(x1 + w_crop, w_img)
                  y1 = max(y2 - h_crop, 0)
                  x1 = max(x2 - w_crop, 0)
                  crop_img = img[:, :, y1:y2, x1:x2]
                  crop_seg_logit = self.encode_decode(crop_img, img_meta)
                  preds += F.pad(crop_seg_logit,
                                 (int(x1), int(preds.shape[3] - x2), int(y1),
                                  int(preds.shape[2] - y2)))
  
                  count_mat[:, :, y1:y2, x1:x2] += 1
          assert (count_mat == 0).sum() == 0
          if torch.onnx.is_in_onnx_export():
              # cast count_mat to constant while exporting to ONNX
              count_mat = torch.from_numpy(
                  count_mat.cpu().detach().numpy()).to(device=img.device)
          preds = preds / count_mat
          if rescale:
              preds = resize(
                  preds,
                  size=img_meta[0]['ori_shape'][:2],
                  mode='bilinear',
                  align_corners=self.align_corners,
                  warning=False)
          return preds
  
      def whole_inference(self, img, img_meta, rescale):
          """Inference with full image."""
  
          seg_logit = self.encode_decode(img, img_meta)
          if rescale:
              seg_logit = resize(
                  seg_logit,
                  size=img_meta[0]['ori_shape'][:2],
                  mode='bilinear',
                  align_corners=self.align_corners,
                  warning=False)
  
          return seg_logit
  
      def inference(self, img, img_meta, rescale,k):
          """Inference with slide/whole style.
  
          Args:
              img (Tensor): The input image of shape (N, 3, H, W).
              img_meta (dict): Image info dict where each dict has: 'img_shape',
                  'scale_factor', 'flip', and may also contain
                  'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
                  For details on the values of these keys see
                  `mmseg/datasets/pipelines/formatting.py:Collect`.
              rescale (bool): Whether rescale back to original shape.
  
          Returns:
              Tensor: The output segmentation map.
          """
  
          assert self.test_cfg.mode in ['slide', 'whole']
          ori_shape = img_meta[0]['ori_shape']
          assert all(_['ori_shape'] == ori_shape for _ in img_meta)
          if self.test_cfg.mode == 'slide':
              seg_logit = self.slide_inference(img, img_meta, rescale)
              # print('1111')
          else:
              seg_logit = self.whole_inference(img, img_meta, rescale)
              # print('\n',seg_logit.size())
          
          if k==7 or k==-1:
              output = F.softmax(seg_logit, dim=1)
          else:
              output = seg_logit
          # print(output.size())
          flip = img_meta[0]['flip']
          if flip:
              flip_direction = img_meta[0]['flip_direction']
              assert flip_direction in ['horizontal', 'vertical']
              if flip_direction == 'horizontal':
                  output = output.flip(dims=(3, ))
              elif flip_direction == 'vertical':
                  output = output.flip(dims=(2, ))
  
          return output
  
      def simple_test(self, img, img_meta, rescale=True, k=7):
          """Simple test with single image."""
  
          # print(k)
          seg_logit = self.inference(img, img_meta, rescale,k)
          # print('\n',seg_logit.size())
          # np.save('embedding',seg_logit)
          
          if k==7 or k==-1:
              seg_pred = seg_logit.argmax(dim=1)
          else:
              seg_pred = kmeans_debackground(img_meta, seg_logit, k+1)[0]
  
          # print(seg_pred.size())
          if torch.onnx.is_in_onnx_export():
              # our inference backend only support 4D output
              seg_pred = seg_pred.unsqueeze(0)
              return seg_pred
          seg_pred = seg_pred.cpu().numpy()
          # unravel batch dim
          seg_pred = list(seg_pred)
          # print(seg_pred)
          return seg_pred
  
      def aug_test(self, imgs, img_metas, rescale=True):
          """Test with augmentations.
  
          Only rescale=True is supported.
          """
          # aug_test rescale all imgs back to ori_shape for now
          assert rescale
          # to save memory, we get augmented seg logit inplace
          seg_logit = self.inference(imgs[0], img_metas[0], rescale)
          for i in range(1, len(imgs)):
              cur_seg_logit = self.inference(imgs[i], img_metas[i], rescale)
              seg_logit += cur_seg_logit
          seg_logit /= len(imgs)
          seg_pred = seg_logit.argmax(dim=1)
          seg_pred = seg_pred.cpu().numpy()
          # unravel batch dim
          seg_pred = list(seg_pred)
          return seg_pred
  