code: def mask_rcnn_inference(
      pred_mask_logits: torch.Tensor, pred_instances: List[Instances]
  ):
      """
      Convert pred_mask_logits to estimated foreground probability masks while also
      extracting only the masks for the predicted classes in pred_instances. For each
      predicted box, the mask of the same class is attached to the instance by adding a
      new "pred_masks" field to pred_instances.
  
      Args:
          pred_mask_logits (Tensor): A tensor of shape (B, C, Hmask, Wmask) or (B, 1, Hmask, Wmask)
              for class-specific or class-agnostic, where B is the total number of predicted masks
              in all images, C is the number of foreground classes, and Hmask, Wmask are the height
              and width of the mask predictions. The values are logits.
          pred_instances (list[Instances]): A list of N Instances, where N is the number of images
              in the batch. Each Instances must have field "pred_classes".
  
      Returns:
          None. pred_instances will contain an extra "pred_masks" field storing a mask of size (Hmask,
              Wmask) for predicted class. Note that the masks are returned as a soft (non-quantized)
              masks the resolution predicted by the network; post-processing steps, such as resizing
              the predicted masks to the original image resolution and/or binarizing them, is left
              to the caller.
      """
      cls_agnostic_mask = pred_mask_logits.size(1) == 1
  
      if cls_agnostic_mask:
          mask_probs_pred = pred_mask_logits.sigmoid()
      else:
          # Select masks corresponding to the predicted classes
          num_masks = pred_mask_logits.shape[0]
          class_pred = cat([i.pred_classes for i in pred_instances])
          device = (
              class_pred.device
              if torch.jit.is_scripting()
              else (
                  "cpu" if torch.jit.is_tracing() else class_pred.device
              )
          )
          indices = move_device_like(
              torch.arange(num_masks, device=device), class_pred
          )
          mask_probs_pred = pred_mask_logits[indices, class_pred][
              :, None
          ].sigmoid()
      # mask_probs_pred.shape: (B, 1, Hmask, Wmask)
  
      num_boxes_per_image = [len(i) for i in pred_instances]
      mask_probs_pred = mask_probs_pred.split(
          num_boxes_per_image, dim=0
      )
  
      for prob, instances in zip(mask_probs_pred, pred_instances):
          instances.pred_masks = prob
  
imports: from typing import List
  import fvcore.nn.weight_init
  import torch
  from torch import nn
  from torch.nn import functional
  from detectron2.config import configurable
  from detectron2.layers import Conv2d, ConvTranspose2d, ShapeSpec, cat, get_norm
  from detectron2.layers.wrappers import move_device_like
  from detectron2.structures import Instances
  from detectron2.utils.events import get_event_storage
  from detectron2.utils.registry import Registry
label_desc: Use image node_id 2 for calling a global function with example usage: mask_rcnn_inference(pred_mask_logits, pred_instances) without return types