diff --git a/deeplabcut/pose_estimation_pytorch/apis/analyze_images.py b/deeplabcut/pose_estimation_pytorch/apis/analyze_images.py index 560b2b85ea..b8142367d9 100644 --- a/deeplabcut/pose_estimation_pytorch/apis/analyze_images.py +++ b/deeplabcut/pose_estimation_pytorch/apis/analyze_images.py @@ -10,6 +10,7 @@ # from __future__ import annotations +import copy import glob import json import logging @@ -21,6 +22,7 @@ import numpy as np from tqdm import tqdm +import deeplabcut.pose_estimation_pytorch.apis.visualization as visualization import deeplabcut.pose_estimation_pytorch.config.utils as config_utils import deeplabcut.pose_estimation_pytorch.modelzoo as modelzoo from deeplabcut.core.engine import Engine @@ -45,67 +47,132 @@ def superanimal_analyze_images( images: str | Path | list[str] | list[Path], max_individuals: int, out_folder: str | Path, - bbox_threshold: float = 0.6, progress_bar: bool = True, device: str | None = None, + pose_threshold: float = 0.4, + bbox_threshold: float = 0.6, + plot_skeleton: bool = True, + customized_model_config: str | Path | dict | None = None, + customized_pose_checkpoint: str | Path | None = None, + customized_detector_checkpoint: str | Path | None = None, ) -> dict[str, dict]: """ - This funciton inferences a superanimal model on a set of images and saves the + This function inferences a superanimal model on a set of images and saves the results as labeled images. Args: - superanimal_name: The name of the superanimal to analyze. Supported list: - - "superanimal_bird" - - "superanimal_topviewmouse" - - "superanimal_quadruped" - model_name: The name of the pose model architecture to use for inference. - detector_name: The name of the detector architecture to use for inference. - images: The images to analyze. Can either be a directory containing images, or + superanimal_name: str + The name of the SuperAnimal to analyze. Supported list: + - "superanimal_bird" + - "superanimal_topviewmouse" + - "superanimal_quadruped" + + model_name: str + The name of the pose model architecture to use for inference. To get a list + of available models for a SuperAnimal, call: + >>> import dlclibrary + >>> superanimal_name = "superanimal_topviewmouse" + >>> dlclibrary.get_available_models(superanimal_name) + + detector_name: str + The name of the detector architecture to use for inference. To get a list + of available detectors for a SuperAnimal, call: + >>> import dlclibrary + >>> superanimal_name = "superanimal_topviewmouse" + >>> dlclibrary.get_available_detectors(superanimal_name) + + images: str, Path, list[str], list[Path] + The images to analyze. Can either be a directory containing images, or a list of paths of images. - max_individuals: The maximum number of individuals to detect in each image. - out_folder: The directory where the labeled images will be saved. - bbox_threshold: The minimum confidence score to keep bounding box detections. - Must be in (0, 1). - progress_bar: Whether to display a progress bar when running inference. - device: The device to use to run image analysis. + + max_individuals: int + The maximum number of individuals to detect in each image. + + out_folder: str | Path + The directory where the labeled images will be saved. + + progress_bar: bool, default=True + Whether to display a progress bar when running inference. + + device: str | None, default=None + The device to use to run image analysis. + + pose_threshold: float, default=0.4 + The cutoff score when plotting pose predictions. To note, this is called + pcutoff in other parts of the code. Must be in (0, 1). + + bbox_threshold: float, default=0.1 + The minimum confidence score to keep bounding box detections. Must be in + (0, 1). + + plot_skeleton: bool, default=True + If a skeleton is defined in the model configuration file, whether to plot + the skeleton connecting the predicted bodyparts on the images. + + customized_model_config: str | Path | dict | None + A customized SuperAnimal model config, as an alternative to the default + SuperAnimal model config. You can get the default SuperAnimal config with: + >>> import deeplabcut.pose_estimation_pytorch.modelzoo as modelzoo + >>> config = modelzoo.load_super_animal_config( + >>> super_animal, model_name, detector_name, + >>> ) + + customized_pose_checkpoint: str | None + A customized SuperAnimal pose checkpoint, as an alternative to the + HuggingFace SuperAnimal models. + + customized_detector_checkpoint: str | None + A customized SuperAnimal detector checkpoint, as an alternative to the + HuggingFace SuperAnimal models. Returns: - The predictions for each image + The predictions made by the model for each image. Examples: - >>> import deeplabcut - >>> from deeplabcut.pose_estimation_pytorch.apis.analyze_images import ( + >>> from deeplabcut.pose_estimation_pytorch.apis import ( >>> superanimal_analyze_images >>> ) - >>> superanimal_name = "superanimal_quadruped" - >>> model_name = "hrnetw32" - >>> device = "cuda:0" - >>> max_individuals = 3 - >>> test_images_folder = "test_rodent_images" - >>> out_images_folder = "vis_test_rodent_images" - >>> ret = superanimal_analyze_images( - >>> superanimal_name, - >>> model_name, - >>> test_images_folder, - >>> max_individuals, - >>> out_images_folder + >>> predictions = superanimal_analyze_images( + >>> superanimal_name="superanimal_topviewmouse", + >>> model_name="resnet_50", + >>> detector_name="fasterrcnn_mobilenet_v3_large_fpn", + >>> images="test_mouse_images", + >>> max_individuals=3, + >>> out_folder="test_mouse_images_labeled", + >>> device="cuda:0", + >>> pose_threshold=0.1, >>> ) """ out_folder = Path(out_folder) out_folder.mkdir(exist_ok=True, parents=True) - snapshot_path = modelzoo.get_super_animal_snapshot_path( - dataset=superanimal_name, model_name=model_name, - ) - detector_path = modelzoo.get_super_animal_snapshot_path( - dataset=superanimal_name, model_name=detector_name, - ) + if customized_pose_checkpoint is None: + snapshot_path = modelzoo.get_super_animal_snapshot_path( + dataset=superanimal_name, + model_name=model_name, + ) + else: + snapshot_path = Path(customized_pose_checkpoint) + + if customized_detector_checkpoint is None: + detector_path = modelzoo.get_super_animal_snapshot_path( + dataset=superanimal_name, + model_name=detector_name, + ) + else: + detector_path = Path(customized_detector_checkpoint) + + if customized_model_config is None: + config = modelzoo.load_super_animal_config( + super_animal=superanimal_name, + model_name=model_name, + detector_name=detector_name, + ) + elif isinstance(customized_model_config, (str, Path)): + config = config_utils.read_config_as_dict(customized_model_config) + else: + config = copy.deepcopy(customized_model_config) - config = modelzoo.load_super_animal_config( - super_animal=superanimal_name, - model_name=model_name, - detector_name=detector_name, - ) config = update_config(config, max_individuals, device) config["metadata"]["individuals"] = [f"animal{i}" for i in range(max_individuals)] if "detector" in config: @@ -121,9 +188,29 @@ def superanimal_analyze_images( progress_bar=progress_bar, ) - superanimal_colormaps = get_superanimal_colormaps() - colormap = superanimal_colormaps[superanimal_name] - create_labeled_images_from_predictions(predictions, out_folder, colormap) + skeleton_bodyparts = config.get("skeleton", []) + skeleton = None + if plot_skeleton and len(skeleton_bodyparts) > 0: + skeleton = [] + bodyparts = config["metadata"]["bodyparts"] + for bpt_0, bpt_1 in skeleton_bodyparts: + skeleton.append( + (bodyparts.index(bpt_0), bodyparts.index(bpt_1)) + ) + + visualization.create_labeled_images( + predictions=predictions, + out_folder=out_folder, + num_bodyparts=len(config["metadata"]["bodyparts"]), + num_unique_bodyparts=len(config["metadata"]["unique_bodyparts"]), + max_individuals=max_individuals, + pcutoff=pose_threshold, + bboxes_pcutoff=bbox_threshold, + cmap=get_superanimal_colormaps()[superanimal_name], + skeleton=skeleton, + skeleton_color=config.get("skeleton_color", "black"), + ) + return predictions @@ -241,7 +328,7 @@ def analyze_image_folder( max_individuals: int | None = None, progress_bar: bool = True, ) -> dict[str, dict[str, np.ndarray | np.ndarray]]: - """Runs pose inference on a folder of images + """Runs pose inference on a folder of images and returns the predictions Args: model_cfg: The model config (or its path) used to analyze the images. @@ -318,32 +405,6 @@ def analyze_image_folder( } -def create_labeled_images_from_predictions(predictions, out_folder, cmap): - for image_path, prediction in predictions.items(): - frame = auxfun_videos.imread(str(image_path), mode="skimage") - fig, ax = plt.subplots() - ax.imshow(frame) - for idx, pose in enumerate(prediction["bodyparts"]): - x, y, confidence = pose[:, 0], pose[:, 1], pose[:, 2] - if np.sum(pose) < 0: - continue - mask = confidence > 0.0 - x = x[mask] - y = y[mask] - ax.scatter(x, y, c=np.arange(len(x)), cmap=cmap) - bboxes = prediction["bboxes"] - for bbox in bboxes: - # Draw bounding boxes around detected objects - xmin, ymin, w, h = bbox - rect = plt.Rectangle( - (xmin, ymin), w, h, fill=False, edgecolor="green", linewidth=2 - ) - - ax.add_patch(rect) - image_name = image_path.split(os.sep)[-1] - fig.savefig(os.path.join(out_folder, f"vis_{image_name}")) - - def plot_images_coco( model_cfg: str | Path | dict, image_folder: str | Path, diff --git a/deeplabcut/pose_estimation_pytorch/apis/visualization.py b/deeplabcut/pose_estimation_pytorch/apis/visualization.py index 72de7c4a13..b4d0cd250b 100644 --- a/deeplabcut/pose_estimation_pytorch/apis/visualization.py +++ b/deeplabcut/pose_estimation_pytorch/apis/visualization.py @@ -14,9 +14,13 @@ from pathlib import Path import cv2 +import matplotlib.collections as collections +import matplotlib.colors as colors +import matplotlib.pyplot as plt import numpy as np import torch import torch.nn.functional as F +from PIL import Image from tqdm import tqdm import deeplabcut.core.visualization as visualization @@ -24,12 +28,143 @@ import deeplabcut.pose_estimation_pytorch.data as data import deeplabcut.pose_estimation_pytorch.data.preprocessor as preprocessor import deeplabcut.pose_estimation_pytorch.models as models +import deeplabcut.utils.visualization as visualization_utils from deeplabcut.core.engine import Engine from deeplabcut.pose_estimation_pytorch.config import read_config_as_dict from deeplabcut.pose_estimation_pytorch.task import Task from deeplabcut.utils import auxiliaryfunctions +def create_labeled_images( + predictions: dict[str, dict[str, np.ndarray | np.ndarray]], + out_folder: str | Path, + num_bodyparts: int, + num_unique_bodyparts: int, + max_individuals: int = 1, + pcutoff: float = 0.6, + bboxes_pcutoff: float = 0.6, + mode: str = "bodypart", + cmap: str | colors.Colormap = "rainbow", + dot_size: int = 12, + alpha_value: float = 0.7, + skeleton: list[tuple[int, int]] | None = None, + skeleton_color: str = "k", +): + """Plots model predictions on images. + + Args: + predictions: The predictions to plot. A dictionary mapping image paths to + the predictions made by the model on that image. The predictions should + contain a "bodyparts" key, mapping to an array of shape (max_individuals, + num_bodyparts, 3) containing predicted bodyparts. If there are any unique + bodyparts predicted, then it should also contain a "unique_bodyparts" key, + mapping to an array of shape (1, num_bodyparts, 3) containing the predicted + unique bodyparts. + out_folder: The folder where model predictions should be saved. + num_bodyparts: The number of bodyparts predicted by the model. + num_unique_bodyparts: The number of unique bodyparts predicted by the model. + max_individuals: The maximum number of individuals predicted by the model. + pcutoff: The p-cutoff score above which predicted bodyparts are displayed with + a "⋅" marker, and below which they are displayed with a "X" marker. + bboxes_pcutoff: The bounding box cutoff score, below which predicted bounding + boxes are shown with a dashed line. + mode: One of "bodypart", "individual". Whether to color predictions by + bodypart or individual. + cmap: The colormap to use to plot predictions. + dot_size: The size of the bodypart prediction markers. + alpha_value: The transparency value of the bodypart prediction markers. + skeleton: If skeletons should be plotted, the list of bodyparts that constitute + the skeletons. + skeleton_color: The color with which to plot the skeleton, if one is given. + """ + out_folder = Path(out_folder) + out_folder.mkdir(exist_ok=True) + + color_by_individual = mode == "individual" + + bboxes_color = "g" + if isinstance(cmap, str): + cmap_name = cmap + num_colors = num_bodyparts + num_unique_bodyparts + 1 + if color_by_individual: + num_colors = max_individuals + 1 + cmap = visualization_utils.get_cmap(num_colors, name=cmap_name) + + bboxes_color = cmap(num_bodyparts + num_unique_bodyparts + 1) + if color_by_individual: + bboxes_color = visualization_utils.get_cmap(num_colors, name=cmap_name) + + fig, ax = visualization_utils.create_minimal_figure() + for image_path, image_predictions in predictions.items(): + # Load frame + frame = Image.open(str(image_path)) + h, w = frame.height, frame.width + + # Get bodypart predictions, put in order so colors are set correctly + pred = image_predictions["bodyparts"] # (num_idv, num_kpt, 3) + bones = None + if skeleton is not None: + bones = [ + idv_pose[[idx_1, idx_2]][:, :2] + for idv_pose in pred + for idx_1, idx_2 in skeleton + ] + if not color_by_individual: + pred = pred.swapaxes(0, 1) + predictions = [p[:, :2] for p in pred] + scores = [p[:, 2:3] for p in pred] + + # Add unique bodypart predictions if there are any + if num_unique_bodyparts > 0: + unique_pred = image_predictions["unique_bodyparts"] + if not color_by_individual: + unique_pred = unique_pred.swapaxes(0, 1) + predictions += [up[:, :2] for up in unique_pred] + scores += [up[:, 2:3] for up in unique_pred] + + # Make empty ground truth as we have none + gt = [np.full((1, 2), fill_value=np.nan) for _ in range(len(predictions))] + + # Load bounding boxes if there are any + bounding_boxes = None + if "bboxes" in image_predictions: + bboxes = image_predictions["bboxes"] + bbox_scores = image_predictions["bbox_scores"] + bounding_boxes = (bboxes, bbox_scores) + + # Create the figure + fig.set_size_inches(w / 100, h / 100) + ax.set_xlim(0, w) + ax.set_ylim(0, h) + ax.invert_yaxis() + visualization_utils.make_multianimal_labeled_image( + np.asarray(frame), + gt, + predictions, + scores, + cmap, + dot_size, + alpha_value, + pcutoff, + ax=ax, + bounding_boxes=bounding_boxes, + bboxes_cutoff=bboxes_pcutoff, + bboxes_color=bboxes_color, + ) + if bones is not None: + ax.add_collection( + collections.LineCollection( + bones, colors=skeleton_color, alpha=alpha_value + ) + ) + + output_path = out_folder / f"predictions_{Path(image_path).stem}.png" + fig.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=0, hspace=0) + fig.savefig(output_path) + visualization_utils.erase_artists(ax) + plt.close(fig) + + @torch.no_grad() def extract_model_outputs( images: list[str] | list[Path], @@ -93,10 +228,7 @@ def extract_model_outputs( output[head]["heatmap"] = F.sigmoid(output[head]["heatmap"]) output = { - head: { - name: output.cpu().numpy() - for name, output in head_outputs.items() - } + head: {name: output.cpu().numpy() for name, output in head_outputs.items()} for head, head_outputs in output.items() } model_data.append( @@ -219,7 +351,8 @@ def extract_maps( snapshot_id = snapshot.path.stem extracted_maps[loader.train_fraction][snapshot_id] = {} runner = utils.get_pose_inference_runner( - model_config=loader.model_cfg, snapshot_path=snapshot.path, + model_config=loader.model_cfg, + snapshot_path=snapshot.path, ) results = extract_model_outputs( image_paths, @@ -254,7 +387,12 @@ def extract_maps( is_train = image_idx in train_idx extracted_maps[loader.train_fraction][snapshot_id][key] = ( - *parsed, None, bpt_names, paf_graph, img_name, is_train + *parsed, + None, + bpt_names, + paf_graph, + img_name, + is_train, ) # img, scmap, locref, paf, peaks, bpt_names, paf_graph, img_name, is_train @@ -405,9 +543,7 @@ def _get_context( bboxes_train = loader.ground_truth_bboxes(mode="train") bboxes_test = loader.ground_truth_bboxes(mode="test") bboxes = {**bboxes_train, **bboxes_test} - return [ - dict(bboxes=bboxes[str(img_path)]) for img_path in image_paths - ] + return [dict(bboxes=bboxes[str(img_path)]) for img_path in image_paths] detector_runner = utils.get_detector_inference_runner( model_config=loader.model_cfg,