class_docstr: Check if the average image brightness (or bbox location if applicable) is similar between train and test set.
  
  The check computes the average grayscale image per dataset (train and test) and compares the resulting images.
  Additionally, in case of an object detection task, the check will compare the average locations of the bounding
  boxes between the datasets.
  
  Parameters
  ----------
  {additional_check_init_params:2*indent}
class_docstr_tok: ["Check","if","the","average","image","brightness","(","or","bbox","location","if","applicable",")","is","similar","between","train","and","test","set",".","The","check","computes","the","average","grayscale","image","per","dataset","(","train","and","test",")","and","compares","the","resulting","images",".","Additionally",",","in","case","of","an","object","detection","task",",","the","check","will","compare","the","average","locations","of","the","bounding","boxes","between","the","datasets",".","Parameters","--","--","--","--","--","{","additional_check_init_params:2","*","indent","}"]
docstr: Create the average images and display them.
  
  Returns
  -------
  CheckResult
      value: The difference images. One for average image brightness, and one for bbox locations if applicable.
      display: Heatmaps for image brightness (train, test, diff) and heatmap for bbox locations if applicable.
docstr_tok: ["Create","the","average","images","and","display","them",".","Returns","--","--","--","-","CheckResult","value",":","The","difference","images",".","One","for","average","image","brightness",",","and","one","for","bbox","locations","if","applicable",".","display",":","Heatmaps","for","image","brightness","(","train",",","test",",","diff",")","and","heatmap","for","bbox","locations","if","applicable","."]
code: def compute(self, context: Context) -> CheckResult:
      """Create the average images and display them.
  
      Returns
      -------
      CheckResult
          value: The difference images. One for average image brightness, and one for bbox locations if applicable.
          display: Heatmaps for image brightness (train, test, diff) and heatmap for bbox locations if applicable.
      """
      # Compute the average grayscale image by dividing the accumulated sum by the number of images
      train_grayscale = (
          np.expand_dims(
              self._grayscale_heatmap[DatasetKind.TRAIN], axis=2
          )
          / self._counter[DatasetKind.TRAIN]
      ).astype(np.uint8)
      test_grayscale = (
          np.expand_dims(
              self._grayscale_heatmap[DatasetKind.TEST], axis=2
          )
          / self._counter[DatasetKind.TEST]
      ).astype(np.uint8)
  
      value = {
          "diff": self._image_diff(test_grayscale, train_grayscale)
      }
  
      if context.with_display:
          # Add a display for the heatmap of the average grayscale image
          display = [
              self.plot_row_of_heatmaps(
                  train_grayscale,
                  test_grayscale,
                  "Compare average image brightness",
                  [context.train.name, context.test.name],
              )
          ]
          display[0].update_layout(
              coloraxis={
                  "colorscale": "Inferno",
                  "cmin": 0,
                  "cmax": 255,
              },
              coloraxis_colorbar={"title": "Pixel Value"},
          )
  
      # If the task is object detection, compute the average heatmap of the bounding box locations by dividing the
      # accumulated sum by the number of images
      if self._task_type == TaskType.OBJECT_DETECTION:
          # bbox image values are frequency, between 0 and 100
          train_bbox = (
              100
              * np.expand_dims(
                  self._bbox_heatmap[DatasetKind.TRAIN], axis=2
              )
              / self._counter[DatasetKind.TRAIN]
              / 255
          ).astype(np.uint8)
          test_bbox = (
              100
              * np.expand_dims(
                  self._bbox_heatmap[DatasetKind.TEST], axis=2
              )
              / self._counter[DatasetKind.TEST]
              / 255
          ).astype(np.uint8)
  
          value["diff_bbox"] = self._image_diff(test_bbox, train_bbox)
  
          if context.with_display:
              display.append(
                  self.plot_row_of_heatmaps(
                      train_bbox,
                      test_bbox,
                      "Compare average label bbox locations",
                      [context.train.name, context.test.name],
                  )
              )
              # bbox image values are frequency, between 0 and 100
              display[1].update_layout(
                  coloraxis={
                      "colorscale": "Inferno",
                      "cmin": 0,
                      "cmax": 100,
                  },
                  coloraxis_colorbar={"title": "% Coverage"},
              )
  
      return CheckResult(
          value=value,
          display=display if context.with_display else None,
          header="Heatmap Comparison",
      )
  
imports: from collections import defaultdict
  from typing import Iterable, List, Optional, Tuple
  import cv2
  import numpy
  import plotly.graph_objs
  from plotly.subplots import make_subplots
  from deepchecks.core import CheckResult, DatasetKind
  from deepchecks.vision._shared_docs import docstrings
  from deepchecks.vision.base_checks import TrainTestCheck
  from deepchecks.vision.context import Context
  from deepchecks.vision.utils.image_functions import apply_heatmap_image_properties, numpy_grayscale_to_heatmap_figure
  from deepchecks.vision.vision_data import TaskType
  from deepchecks.vision.vision_data.batch_wrapper import BatchWrapper
label_desc: Use image node_id 4 for calling the HeatmapComparison obj's underlying member method code with example usage: obj.compute(context) and returns: CheckResult