code: def merge_aug_masks(
      aug_masks: List[Tensor],
      img_metas: dict,
      weights: Optional[Union[list, Tensor]] = None,
  ) -> Tensor:
      """Merge augmented mask prediction.
  
      Args:
          aug_masks (list[Tensor]): each has shape
              (n, c, h, w).
          img_metas (dict): Image information.
          weights (list or Tensor): Weight of each aug_masks,
              the length should be n.
  
      Returns:
          Tensor: has shape (n, c, h, w)
      """
      recovered_masks = []
      for i, mask in enumerate(aug_masks):
          if weights is not None:
              assert len(weights) == len(aug_masks)
              weight = weights[i]
          else:
              weight = 1
          flip = img_metas.get("flip", False)
          if flip:
              flip_direction = img_metas["flip_direction"]
              if flip_direction == "horizontal":
                  mask = mask[:, :, :, ::-1]
              elif flip_direction == "vertical":
                  mask = mask[:, :, ::-1, :]
              elif flip_direction == "diagonal":
                  mask = mask[:, :, :, ::-1]
                  mask = mask[:, :, ::-1, :]
              else:
                  raise ValueError(
                      f"Invalid flipping direction '{flip_direction}'"
                  )
          recovered_masks.append(mask[None, :] * weight)
  
      merged_masks = torch.cat(recovered_masks, 0).mean(dim=0)
      if weights is not None:
          merged_masks = merged_masks * len(weights) / sum(weights)
      return merged_masks
  
file_path: mmdetection/mmdet/models/test_time_augs/merge_augs.py
imports: import copy
  import warnings
  from typing import List, Optional, Union
  import numpy
  import torch
  from mmcv.ops import nms
  from mmengine.config import ConfigDict
  from torch import Tensor
  from mmdet.structures.bbox import bbox_mapping_back
label_desc: Use image node_id 5 for calling a global function with example usage: merge_aug_masks(aug_masks, img_metas, weights) and returns: merged_masks