code: def prepare_coco_eval_dicts(
      batch: Dict[str, Any],
      predictions: Dict[str, jnp.ndarray],
      logits_to_probs_fn: Callable[[jnp.ndarray], jnp.ndarray],
      gather: bool = False,
  ) -> Tuple[Dict[str, jnp.ndarray], Dict[str, jnp.ndarray]]:
      """Prepare predictions and validations for COCO eval.
  
      Args:
        batch: Eval batch targets.
        predictions: Predictions from DeformableDETR.
        logits_to_probs_fn: Convert logits to probabilities.
        gather: Whether to perform gather if we are in a pmapped eval.
  
      Returns:
        Targets and predictions formatted for COCO eval.
      """
      pred_probs = logits_to_probs_fn(predictions["pred_logits"])
      # Collect necessary predictions and target information from all hosts.
      predictions_out = {
          "pred_probs": pred_probs,
          "pred_logits": predictions["pred_logits"],
          "pred_boxes": predictions["pred_boxes"],
      }
      labels = {
          "image/id": batch["label"]["image/id"],
          "size": batch["label"]["size"],
          "orig_size": batch["label"]["orig_size"],
      }
      to_copy = [
          "labels",
          "boxes",
          "not_exhaustive_category_ids",
          "neg_category_ids",
      ]
      for name in to_copy:
          if name in batch["label"]:
              labels[name] = batch["label"][name]
  
      targets = {"label": labels, "batch_mask": batch["batch_mask"]}
  
      if gather:
          predictions_out = jax.lax.all_gather(predictions_out, "batch")
          targets = jax.lax.all_gather(targets, "batch")
      return targets, predictions_out
  
file_path: scenic/scenic/projects/baselines/deformable_detr/coco_eval.py
imports: import copy
  from typing import Any, Callable, Dict, Optional, Tuple
  import jax
  import jax.numpy
  import numpy
  from scenic.dataset_lib.coco_dataset import coco_eval
  from scenic.model_lib.base_models import box_utils
  from scenic.projects.baselines.detr.train_utils import DetrGlobalEvaluator
  import scipy
label_desc: Use image node_id 1 for calling a global function with example usage: prepare_coco_eval_dicts(batch, predictions, logits_to_probs_fn, gather) and returns: targets,  predictions_out