Repository: Project-MONAI/MONAI
File: monai/utils/jupyter_utils.py (Python)

Code before:
        im.axis("off")
        axes.append(im)

    return axes


def tensor_to_images(name: str, tensor: torch.Tensor) -> np.ndarray | None:
    """
    Return an tuple of images derived from the given tensor. The `name` value indices which key from the
    output or batch value the tensor was stored as, or is "Batch" or "Output" if these were single tensors
    instead of dictionaries. Returns a tuple of 2D images of shape HW, or 3D images of shape CHW where C is
    color channels RGB or RGBA. This allows multiple images to be created from a single tensor, ie. to show
    each channel separately.
    """
    if tensor.ndim == 3 and tensor.shape[1] > 2 and tensor.shape[2] > 2:
        return tensor.cpu().data.numpy()  # type: ignore[no-any-return]
    if tensor.ndim == 4 and tensor.shape[2] > 2 and tensor.shape[3] > 2:
        dmid = tensor.shape[1] // 2
        return tensor[:, dmid].cpu().data.numpy()  # type: ignore[no-any-return]

    return None


def plot_engine_status(
    engine: Engine,
    logger: MetricLogger,
    title: str = "Training Log",
    yscale: str = "log",
    avg_keys: tuple[str] = (LOSS_NAME,),
    window_fraction: int = 20,
    image_fn: Callable[[str, torch.Tensor], Any] | None = tensor_to_images,
    fig: plt.Figure = None,
    selected_inst: int = 0,
) -> tuple[plt.Figure, list]:
    """
    Plot the status of the given Engine with its logger. The plot will consist of a graph of loss values and metrics
    taken from the logger, and images taken from the `output` and `batch` members of `engine.state`. The images are
    converted to Numpy arrays suitable for input to `Axes.imshow` using `image_fn`, if this is None then no image
    plotting is done.

    Args:
        engine: Engine to extract images from
        logger: MetricLogger to extract loss and metric data from
        title: graph title
        yscale: for metric plot, scale for y-axis compatible with `Axes.set_yscale`
        avg_keys: for metric plot, tuple of keys in `graphmap` to provide running average plots for
        window_fraction: for metric plot, what fraction of the graph value length to use as the running average window
        image_fn: callable converting tensors keyed to a name in the Engine to a tuple of images to plot
        fig: Figure object to plot into, reuse from previous plotting for flicker-free refreshing
        selected_inst: index of the instance to show in the image plot


Review comment (suggestion):
Building the docs doesn't work because of this forward reference to `MetricLogger`. I tried putting it into quotes so `logger: "MetricLogger"` but this is automatically reverted by the pre-commit pyupgrade. I also searched for inline comments to locally disable pyupgrade, but this does not exist.

So my solutions would be one of:
- Completely disable pyupgrade for this file
- (Preferred) Use `Any` for the argument

Code after:

    return axes


def tensor_to_images(name: str, tensor: torch.Tensor) -> np.ndarray | None:
    """
    Return an tuple of images derived from the given tensor. The `name` value indices which key from the
    output or batch value the tensor was stored as, or is "Batch" or "Output" if these were single tensors
    instead of dictionaries. Returns a tuple of 2D images of shape HW, or 3D images of shape CHW where C is
    color channels RGB or RGBA. This allows multiple images to be created from a single tensor, ie. to show
    each channel separately.
    """
    if tensor.ndim == 3 and tensor.shape[1] > 2 and tensor.shape[2] > 2:
        return tensor.cpu().data.numpy()  # type: ignore[no-any-return]
    if tensor.ndim == 4 and tensor.shape[2] > 2 and tensor.shape[3] > 2:
        dmid = tensor.shape[1] // 2
        return tensor[:, dmid].cpu().data.numpy()  # type: ignore[no-any-return]

    return None


def plot_engine_status(
    engine: Engine,
    logger: Any,
    title: str = "Training Log",
    yscale: str = "log",
    avg_keys: tuple[str] = (LOSS_NAME,),
    window_fraction: int = 20,
    image_fn: Callable[[str, torch.Tensor], Any] | None = tensor_to_images,
    fig: plt.Figure = None,
    selected_inst: int = 0,
) -> tuple[plt.Figure, list]:
    """
    Plot the status of the given Engine with its logger. The plot will consist of a graph of loss values and metrics
    taken from the logger, and images taken from the `output` and `batch` members of `engine.state`. The images are
    converted to Numpy arrays suitable for input to `Axes.imshow` using `image_fn`, if this is None then no image
    plotting is done.

    Args:
        engine: Engine to extract images from
        logger: MetricLogger to extract loss and metric data from
        title: graph title
        yscale: for metric plot, scale for y-axis compatible with `Axes.set_yscale`
        avg_keys: for metric plot, tuple of keys in `graphmap` to provide running average plots for
        window_fraction: for metric plot, what fraction of the graph value length to use as the running average window
        image_fn: callable converting tensors keyed to a name in the Engine to a tuple of images to plot
        fig: Figure object to plot into, reuse from previous plotting for flicker-free refreshing
        selected_inst: index of the instance to show in the image plot

    Returns:
        Figure object (or `fig` if given), list of Axes objects for graph and images