Function generate_model_output_plots

deeplabcut/core/visualization.py:143–234

Generates model output plots (maps) for an image and saves them to disk. Args: output_folder: The folder in which the plots should be saved. image_name: The name of the image for which the plots were generated. bodypart_names: The names of bodyparts the model outputs.

(
    output_folder: Path,
    image_name: str,
    bodypart_names: list[str],
    bodyparts_to_plot: list[str],
    image: np.ndarray,
    scmap: np.ndarray,
    locref: np.ndarray | None = None,
    paf: np.ndarray | None = None,
    paf_graph: list[tuple[int, int]] | None = None,
    paf_all_in_one: bool = True,
    paf_colormap: str = "rainbow",
    output_suffix: str = "",
)

def generate_model_output_plots(
    output_folder: Path,
    image_name: str,
    bodypart_names: list[str],
    bodyparts_to_plot: list[str],
    image: np.ndarray,
    scmap: np.ndarray,
    locref: np.ndarray | None = None,
    paf: np.ndarray | None = None,
    paf_graph: list[tuple[int, int]] | None = None,
    paf_all_in_one: bool = True,
    paf_colormap: str = "rainbow",
    output_suffix: str = "",
) -> None:
    """Generates model output plots (maps) for an image and saves them to disk.

    Args:
        output_folder: The folder in which the plots should be saved.
        image_name: The name of the image for which the plots were generated.
        bodypart_names: The names of bodyparts the model outputs.
        bodyparts_to_plot: The names of bodyparts that should be plot.
        image: Shape (height, width, channels). The image on which the model was run.
        scmap: Shape (height, width, num_bodyparts). The scoremaps output by the model.
        locref: Shape (height, width, num_bodyparts, 2). Optionally, the location
            refinement fields output by the model.
        paf: Shape (height, width, 2 * len(paf_graph)). Optionally, the part-affinity
            fields output by the model.
        paf_graph: Must be set if paf is not None. The PAF graph used to assemble.
        paf_all_in_one: Whether to plot all PAFs in a single image.
        paf_colormap: The colormap to use for the PAF maps.
        output_suffix: The filename suffix for the maps to output.
    """

    def _filename(map_name) -> str:
        return f"{image_name}_{map_name}_{output_suffix}.png"

    to_plot = [i for i, bpt in enumerate(bodypart_names) if bpt in bodyparts_to_plot]
    if len(to_plot) > 1:
        map_ = scmap[:, :, to_plot].sum(axis=2)
    elif len(to_plot) == 1 and len(bodypart_names) > 1:
        map_ = scmap[:, :, to_plot[0]]
    else:
        map_ = scmap[..., 0]

    fig1, _ = visualize_scoremaps(image, map_)
    fig1.savefig(output_folder / _filename("scmap"))

    if locref is not None:
        if len(to_plot) > 1:
            map_ = scmap[:, :, to_plot]
            locref_x_ = locref[:, :, to_plot, 0]
            locref_y_ = locref[:, :, to_plot, 1]
            # only get the locref fields around their respective detections
            locref_x_[map_ < 0.5] = 0
            locref_y_[map_ < 0.5] = 0
            # combine locrefs
            map_ = map_.sum(axis=2)
            locref_x_ = locref_x_.sum(axis=2)