diff --git a/deeplabcut/modelzoo/generalized_data_converter/datasets/base.py b/deeplabcut/modelzoo/generalized_data_converter/datasets/base.py index 77f5b635a4..8ea478ce00 100644 --- a/deeplabcut/modelzoo/generalized_data_converter/datasets/base.py +++ b/deeplabcut/modelzoo/generalized_data_converter/datasets/base.py @@ -153,11 +153,20 @@ def populate_generic(self): raise NotImplementedError("Must implement this function") def materialize( - self, proj_root, framework="coco", deepcopy=False, append_image_id=True + self, + proj_root, + framework="coco", + deepcopy=False, + append_image_id=True, + no_image_copy=False, ): mat_func = mat_func_factory(framework) self.meta["mat_datasets"] = {self.meta["dataset_name"]: self} self.meta["imageid2datasetname"] = self.imageid2datasetname + kwargs = dict(deepcopy=deepcopy, append_image_id=append_image_id) + if framework == "coco": + kwargs["no_image_copy"] = no_image_copy + mat_func( proj_root, self.generic_train_images, @@ -165,8 +174,7 @@ def materialize( self.generic_train_annotations, self.generic_test_annotations, self.meta, - deepcopy=deepcopy, - append_image_id=append_image_id, + **kwargs, ) def whether_anno_image_match(self, images, annotations): diff --git a/deeplabcut/modelzoo/generalized_data_converter/datasets/materialize.py b/deeplabcut/modelzoo/generalized_data_converter/datasets/materialize.py index dd5224a7a1..63211b8f35 100644 --- a/deeplabcut/modelzoo/generalized_data_converter/datasets/materialize.py +++ b/deeplabcut/modelzoo/generalized_data_converter/datasets/materialize.py @@ -12,6 +12,7 @@ import os import pickle import shutil +from pathlib import Path import numpy as np import pandas as pd @@ -660,9 +661,10 @@ def _generic2coco( train_annotations, test_annotations, meta, - deepcopy=False, - full_image_path=True, - append_image_id=True, + deepcopy: bool = False, + full_image_path: bool = True, + append_image_id: bool = True, + no_image_copy: bool = False, ): """ Take generic data and create coco structure @@ -672,6 +674,17 @@ def _generic2coco( annotations - train.json - test.json + + Args: + deepcopy: Only when no_image_copy=False. If False, images are not copied from + their original location and symlinks are created instead. + full_image_path: Only when no_image_copy=False. If True, the ``file_name`` for + the images in the annotation files contain the resolved path to the images. + Otherwise, a relative path is used. + append_image_id: Only when no_image_copy=False. Appends the image IDs in the + dataset to the image names. + no_image_copy: Instead of copying images to the COCO dataset, the full paths to + the images in the original dataset are used in the annotations. """ os.makedirs(os.path.join(proj_root, "images"), exist_ok=True) @@ -693,54 +706,46 @@ def _generic2coco( broken_links = [] # copying images via symbolic link for image in train_images + test_images: - src = image["file_name"] + # important to resolve the filepath! Otherwise, errors can occur when running + # this code from Jupyter Notebooks + src = Path(image["file_name"]).resolve() image_id = image["id"] - if not os.path.exists(src): + if not src.exists(): print("problem comes from", image["source_dataset"]) print(src) broken_links.append(image_id) continue - else: - pass - # print ('success comes from', image['source_dataset']) - # print (src) - - # in dlc, some images have same name but under different folder - # we used to use a parent folder to distinguish them, but it's only applicable to DLC - # so here it's easier to just append a id into the filename - image_name = src.split(os.sep)[-1] + file_name = str(src) + dest = src + if not no_image_copy: + # in dlc, some images have same name but under different folder + # we used to use a parent folder to distinguish them, but it's only + # applicable to DLC so here it's easier to append an id into the filename - if image_name.count(".") > 1: - sep = image_name.rfind(".") - pre, suffix = image_name[:sep], image_name[sep + 1 :] - else: - # this does not work for image file that looks like image9.5.jpg.. - pre, suffix = image_name.split(".") - - # not to repeatedly add image id in memory replay training - if append_image_id: - dest_image_name = f"{pre}_{image_id}.{suffix}" - else: - dest_image_name = image_name - dest = os.path.join(proj_root, "images", dest_image_name) + # not to repeatedly add image id in memory replay training + dest_image_name = src.name + if append_image_id: + dest_image_name = f"{src.stem}_{image_id}{src.suffix}" - # now, we will also need to update the path in the config files + dest = Path(proj_root) / "images" / dest_image_name + dest = dest.resolve() - if full_image_path: - image["file_name"] = dest - else: - image["file_name"] = os.path.join("images", dest_image_name) + file_name = str(Path(*dest.parts[-2:])) + if full_image_path: + file_name = str(dest) - if deepcopy: - shutil.copy(src, dest) - else: - try: - os.symlink(src, dest) - except: - pass + if deepcopy: + shutil.copy(src, dest) + else: + try: + os.symlink(src, dest) + except Exception as err: + print(f"Could not create a symlink from {src} to {dest}: {err}") + pass + image["file_name"] = file_name lookuptable[dest] = src train_annotations = [ diff --git a/deeplabcut/pose_estimation_pytorch/apis/__init__.py b/deeplabcut/pose_estimation_pytorch/apis/__init__.py index 1347cfaa14..8bdaa0632a 100644 --- a/deeplabcut/pose_estimation_pytorch/apis/__init__.py +++ b/deeplabcut/pose_estimation_pytorch/apis/__init__.py @@ -9,7 +9,10 @@ # Licensed under GNU Lesser General Public License v3.0 # -from deeplabcut.pose_estimation_pytorch.apis.analyze_images import analyze_images +from deeplabcut.pose_estimation_pytorch.apis.analyze_images import ( + analyze_images, + superanimal_analyze_images, +) from deeplabcut.pose_estimation_pytorch.apis.analyze_videos import analyze_videos from deeplabcut.pose_estimation_pytorch.apis.convert_detections_to_tracklets import ( convert_detections2tracklets, diff --git a/deeplabcut/pose_estimation_pytorch/apis/analyze_images.py b/deeplabcut/pose_estimation_pytorch/apis/analyze_images.py index 6263cf3759..26bcfa3fa0 100644 --- a/deeplabcut/pose_estimation_pytorch/apis/analyze_images.py +++ b/deeplabcut/pose_estimation_pytorch/apis/analyze_images.py @@ -47,6 +47,7 @@ def superanimal_analyze_images( images: str | Path | list[str] | list[Path], max_individuals: int, out_folder: str, + bbox_threshold: float = 0.6, progress_bar: bool = True, device: str | None = None, customized_pose_checkpoint: str | None = None, @@ -54,19 +55,18 @@ def superanimal_analyze_images( customized_model_config: str | None = None, ): """ - This funciton inferences a superanimal model on a set of images and saves the results as labeled images. + This funciton inferences a superanimal model on a set of images and saves the + results as labeled images. Parameters ---------- superanimal_name: str - The name of the superanimal to analyze. - supported list: - superanimal_topviewmouse - superanimal_quadruped + The name of the superanimal to analyze. Supported list: + - "superanimal_topviewmouse" + - "superanimal_quadruped" model_name: str - The name of the model to use for inference. - supported list: - hrnetw32 + The name of the model to use for inference. Supported list: + - "hrnetw32" 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. @@ -74,6 +74,10 @@ def superanimal_analyze_images( The maximum number of individuals to detect in each image. out_folder: str The directory where the labeled images will be saved. + bbox_threshold: float, default=0.1 + The minimum confidence score to keep bounding box detections. Must be in (0, 1). + Only used when `customized_model_config=None` (otherwise, edit your + `customized_model_config` with the desired bbox_threshold). progress_bar: bool Whether to display a progress bar when running inference. device: str | None @@ -95,17 +99,19 @@ def superanimal_analyze_images( -------- >>> import deeplabcut >>> from deeplabcut.pose_estimation_pytorch.apis.analyze_images import superanimal_analyze_images - >>> superanimal_name = 'superanimal_quadruped' - >>> model_name = 'hrnetw32' - >>> device = 'cuda' + >>> superanimal_name = "superanimal_quadruped" + >>> model_name = "hrnetw32" + >>> device = "cuda" >>> 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) + >>> 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 + >>> ) """ os.makedirs(out_folder, exist_ok=True) @@ -119,6 +125,10 @@ def superanimal_analyze_images( snapshot_path, detector_path, ) = get_config_model_paths(superanimal_name, model_name) + + if "detector" in model_cfg: + model_cfg["detector"]["model"]["box_score_thresh"] = bbox_threshold + config = {**project_config, **model_cfg} config = update_config(config, max_individuals, device) else: @@ -146,9 +156,7 @@ def superanimal_analyze_images( superanimal_colormaps = get_superanimal_colormaps() colormap = superanimal_colormaps[superanimal_name] - create_labeled_images_from_predictions(predictions, out_folder, colormap) - return predictions @@ -164,8 +172,6 @@ def analyze_images( device: str | None = None, max_individuals: int | None = None, progress_bar: bool = True, - superanimal_name=None, - model_name=None, ) -> dict[str, dict]: """Runs analysis on images using a pose model. diff --git a/deeplabcut/pose_estimation_pytorch/apis/train.py b/deeplabcut/pose_estimation_pytorch/apis/train.py index 4a006d3f09..efec0b6db0 100644 --- a/deeplabcut/pose_estimation_pytorch/apis/train.py +++ b/deeplabcut/pose_estimation_pytorch/apis/train.py @@ -257,8 +257,10 @@ def train_network( dataset_params = loader.get_dataset_parameters() backbone_name = loader.model_cfg["model"]["backbone"]["model_name"] model_name = modelzoo_utils.get_pose_model_type(backbone_name) - # at some point train_network should support a different train_file passing so memory replay can also take the same train file + # at some point train_network should support a different train_file passing + # so memory replay can also take the same train file + print("Preparing data for memory replay (this can take some time)") prepare_memory_replay( loader.project_path, shuffle, @@ -271,6 +273,7 @@ def train_network( customized_pose_checkpoint=weight_init.customized_pose_checkpoint, ) + print("Loading memory replay data") loader = COCOLoader( project_root=Path(loader.model_folder).parent / "memory_replay", model_config_path=loader.model_config_path, diff --git a/deeplabcut/pose_estimation_pytorch/config/base/base.yaml b/deeplabcut/pose_estimation_pytorch/config/base/base.yaml index a507625aaf..5121ee9d97 100644 --- a/deeplabcut/pose_estimation_pytorch/config/base/base.yaml +++ b/deeplabcut/pose_estimation_pytorch/config/base/base.yaml @@ -5,7 +5,7 @@ runner: gpus: null key_metric: "test.mAP" key_metric_asc: true - eval_interval: 1 + eval_interval: 10 optimizer: type: AdamW params: diff --git a/deeplabcut/pose_estimation_pytorch/modelzoo/memory_replay.py b/deeplabcut/pose_estimation_pytorch/modelzoo/memory_replay.py index f7e547f705..1432c2d1ef 100644 --- a/deeplabcut/pose_estimation_pytorch/modelzoo/memory_replay.py +++ b/deeplabcut/pose_estimation_pytorch/modelzoo/memory_replay.py @@ -19,7 +19,6 @@ import numpy as np from scipy.optimize import linear_sum_assignment from scipy.spatial import distance -from scipy.spatial.distance import cdist import deeplabcut.utils.auxiliaryfunctions as af from deeplabcut.core.engine import Engine @@ -27,131 +26,210 @@ from deeplabcut.modelzoo.generalized_data_converter.datasets import ( COCOPoseDataset, MaDLCPoseDataset, - MultiSourceDataset, SingleDLCPoseDataset, ) from deeplabcut.pose_estimation_pytorch.apis.utils import get_inference_runners -from deeplabcut.pose_estimation_pytorch.modelzoo.utils import ( - get_config_model_paths, - update_config, -) -from deeplabcut.utils.pseudo_label import calculate_iou, optimal_match, xywh2xyxy +import deeplabcut.pose_estimation_pytorch.config.utils as config_utils +from deeplabcut.pose_estimation_pytorch.modelzoo.utils import get_config_model_paths +from deeplabcut.pose_estimation_pytorch.runners import InferenceRunner +from deeplabcut.utils.pseudo_label import calculate_iou -# this is reading from a coco project -def prepare_memory_replay_dataset( - source_dataset_folder, - superanimal_name, - model_name, - max_individuals=1, - train_file="train.json", - test_file="test.json", - pose_threshold=0.0, - device=None, - pose_model_path="", - detector_path="", - customized_pose_checkpoint=None, -): - """ - Need to first run inference on the source project train file +def get_superanimal_inference_runners( + superanimal_name: str, + model_name: str, + max_individuals: int, + device: str | None = None, +) -> tuple[InferenceRunner, InferenceRunner | None]: + """Creates the inference runners for SuperAnimal models + + Args: + superanimal_name: the name of the SuperAnimal dataset for which to load runners + model_name: the name of the model + max_individuals: the maximum number of individuals to detect per image + device: the device on which to run + + Returns: + the pose runner for the SuperAnimal model + the detector runner for the SuperAnimal model, if it's a top-down model """ - ( model_config, project_config, pose_model_path, detector_path, ) = get_config_model_paths(superanimal_name, model_name) - - if customized_pose_checkpoint is not None: - print( - "memory replay fine-tuning pose checkpoint is replaced by", - customized_pose_checkpoint, - ) - - config = {**project_config, **model_config} - config = update_config(config, max_individuals, device) - individuals = [f"animal{i}" for i in range(max_individuals)] - config["individuals"] = individuals - num_bodyparts = len(config["bodyparts"]) - train_file_path = os.path.join(source_dataset_folder, "annotations", train_file) - - pose_runner, detector_runner = get_inference_runners( - config, + model_config["metadata"]["individuals"] = [ + f"animal{i}" for i in range(max_individuals) + ] + model_config = config_utils.replace_default_values( + model_config, + num_bodyparts=len(model_config["metadata"]["bodyparts"]), + num_individuals=max_individuals, + backbone_output_channels=model_config["model"]["backbone_output_channels"], + ) + return get_inference_runners( + model_config, snapshot_path=pose_model_path, max_individuals=max_individuals, num_bodyparts=len(model_config["metadata"]["bodyparts"]), num_unique_bodyparts=0, + device=device, detector_path=detector_path, ) - with open(train_file_path, "r") as f: - train_obj = json.load(f) - images = train_obj["images"] - annotations = train_obj["annotations"] - categories = train_obj["categories"] - imagename2id = {} - imageid2name = {} - imagename2gt = defaultdict(list) - - for image in images: - # this only works with relative path as the testing image can be at a different folder - imagename = image["file_name"].split(os.sep)[-1] - imagename2id[imagename] = image["id"] - imageid2name[image["id"]] = imagename +def get_pose_predictions( + project_root: Path, + images: list[str], + bboxes: dict[str, list], + superanimal_name: str, + model_name: str, + max_individuals: int, + device: str | None = None, +) -> dict[str, dict]: + """Gets predictions made by a SuperAnimal model on a DeepLabCut project + + Args: + project_root: The path to the root of the project. + images: The images on which to run inference with the SuperAnimal model. + bboxes: The ground truth bounding boxes for each image in the project. + superanimal_name: The name of the SuperAnimal dataset to use. + model_name: The name of the model to use. + max_individuals: The maximum number of individuals to detect per image. + device: The CUDA device to use. + + Returns: + The predictions made by the SuperAnimal model on each image in the images list. + """ + predictions_folder = project_root / "memory_replay" / superanimal_name / model_name + predictions_folder.mkdir(exist_ok=True, parents=True) + predictions_file = predictions_folder / "pseudo-labels.json" + + # COCO-format annotations file containing predictions made by the SuperAnimal model + sa_predictions = {} + if predictions_file.exists(): + with open(predictions_file, "r") as f: + raw_sa_predictions = json.load(f) + + # parse predictions to convert lists to numpy arrays + for image, predictions in raw_sa_predictions.items(): + sa_predictions[image] = { + "bodyparts": np.array(predictions["bodyparts"]), + "bboxes": np.array(predictions["bboxes"]), + # "bbox_scores": np.array(predictions["bbox_scores"]), + } + + # get images that need to be processed + processed_images = set(sa_predictions.keys()) + images_to_process = [image for image in (set(images) - processed_images)] + + # if all images have been processed by the SuperAnimal model, return the predictions + if len(images_to_process) == 0: + return sa_predictions + + pose_runner, detector_runner = get_superanimal_inference_runners( + superanimal_name, + model_name, + max_individuals, + device=device, + ) - imagename2bbox = defaultdict(list) - for anno in annotations: - imagename = imageid2name[anno["image_id"]] - imagename2gt[imagename].append(anno) - imagename2bbox[imagename].append(anno["bbox"]) + # FIXME(niels, yeshaokai) - Use the detector to combine GT-keypoint created bounding + # boxes and predicted bounding boxes - keep the larger of the two + # bbox_predictions = detector_runner.inference(images=images_to_process) + pose_inputs = [ + ( + project_root / Path(image), + {"bboxes": np.array(bboxes[image])} + ) + for image in images_to_process + ] + predictions = pose_runner.inference(pose_inputs) - imageid2annotations = defaultdict(list) + for image, prediction in zip(images_to_process, predictions): + sa_predictions[image] = prediction - imageids = list(imagename2id.values()) - for annotation in annotations: - image_id = annotation["image_id"] - if annotation["image_id"] in imageids: - imageid2annotations[image_id].append(annotation) + # save the updated SuperAnimal predictions + json_sa_predictions = { + image: { + "bodyparts": predictions["bodyparts"].tolist(), + "bboxes": predictions["bboxes"].tolist(), + # "bbox_scores": predictions["bbox_scores"].tolist(), + } + for image, predictions in sa_predictions.items() + } + with open(predictions_file, "w") as f: + json.dump(json_sa_predictions, f, indent=2) - # need to support more image types - image_extensions = ["*.png", "*.jpg", "*.jpeg", "*.bmp", "*.gif", "*.tiff"] + return sa_predictions - images_in_folder = [] - for ext in image_extensions: - images_in_folder.extend( - glob.glob(os.path.join(source_dataset_folder, "images", ext)) - ) - corresponded_images = [] - for image in images_in_folder: - image_path = image - imagename = image.split(os.sep)[-1] - if imagename in imagename2id: - corresponded_images.append(image_path) +# this is reading from a coco project +def prepare_memory_replay_dataset( + project_root: str | Path, + source_dataset_folder: str | Path, + superanimal_name: str, + model_name: str, + max_individuals: int = 1, + train_file: str = "train.json", + pose_threshold: float = 0.0, + device: str | None = None, + customized_pose_checkpoint: str | None = None, +): + """ + Need to first run inference on the source project train file + """ + project_root = Path(project_root).resolve() + source_dataset_folder = Path(source_dataset_folder).resolve() - images = corresponded_images + if customized_pose_checkpoint is not None: + print( + "memory replay fine-tuning pose checkpoint is replaced by", + customized_pose_checkpoint, + ) - bbox_predictions = detector_runner.inference(images=images) + # Contains the ground truth annotations for the DeepLabCut project + # .../dlc-models-pytorch/.../...shuffle0/train/memory_replay/annotations/train.json + with open(source_dataset_folder / "annotations" / train_file, "r") as f: + project_gt = json.load(f) - bbox_gts = [ - {"bboxes": np.array(imagename2bbox[image.split(os.sep)[-1]])} - for image in images - ] + # parse the GT so that image paths are in the format (no matter the OS): + # "labeled-data/{video_name}/{image_name}" + for image in project_gt["images"]: + image["file_name"] = "/".join(Path(image["file_name"]).parts[-3:]) - pose_inputs = list(zip(images, bbox_gts)) + image_id_to_name = {} + image_id_to_annotations = defaultdict(list) - # pose inference should return meta data for pseudo labeling - predictions = pose_runner.inference(pose_inputs) + image_name_to_id = {} + image_name_to_gt = defaultdict(list) + image_name_to_bbox = defaultdict(list) - assert len(images) == len(predictions) + for image in project_gt["images"]: + image_name_to_id[image["file_name"]] = image["id"] + image_id_to_name[image["id"]] = image["file_name"] - imagename2prediction = {} + for anno in project_gt["annotations"]: + name = image_id_to_name[anno["image_id"]] + image_name_to_gt[name].append(anno) + image_name_to_bbox[name].append(anno["bbox"]) - for image_path, prediction in zip(images, predictions): - imagename = image_path.split(os.sep)[-1] - imagename2prediction[imagename] = prediction + image_ids = list(image_name_to_id.values()) + for annotation in project_gt["annotations"]: + image_id = annotation["image_id"] + if annotation["image_id"] in image_ids: + image_id_to_annotations[image_id].append(annotation) + + image_name_to_prediction = get_pose_predictions( + project_root=project_root, + images=[image["file_name"] for image in project_gt["images"]], + bboxes=image_name_to_bbox, + superanimal_name=superanimal_name, + model_name=model_name, + max_individuals=max_individuals, + device=device, + ) def xywh2xyxy(bbox): temp_bbox = np.copy(bbox) @@ -173,10 +251,11 @@ def optimal_match(gts_list, preds_list): return col_ind - for imagename, gts in imagename2gt.items(): + num_bodyparts = len(project_gt["categories"][0]["keypoints"]) + for image_name, gts in image_name_to_gt.items(): bbox_gts = [np.array(gt["bbox"]) for gt in gts] bbox_gts = [xywh2xyxy(e) for e in bbox_gts] - prediction = imagename2prediction[imagename] + prediction = image_name_to_prediction[image_name] bbox_preds = [xywh2xyxy(pred) for pred in prediction["bboxes"]] optimal_pred_indices = optimal_match(bbox_gts, bbox_preds) @@ -203,10 +282,7 @@ def optimal_match(gts_list, preds_list): # after the mixing, we don't care about confidence anymore for kpt_idx in range(len(matched_gt)): - if ( - matched_gt[kpt_idx][2] < pose_threshold - and matched_gt[kpt_idx][2] > 0 - ): + if 0 < matched_gt[kpt_idx][2] < pose_threshold: matched_gt[kpt_idx][2] = -1 elif matched_gt[kpt_idx][2] > 0: matched_gt[kpt_idx][2] = 2 @@ -218,8 +294,13 @@ def optimal_match(gts_list, preds_list): source_dataset_folder, "annotations", "memory_replay_train.json" ) + # parse the GT to put the image paths back into OS-specific format + for image in project_gt["images"]: + image_rel_path = image["file_name"].split("/") + image["file_name"] = str(project_root.resolve() / Path(*image_rel_path)) + with open(memory_replay_train_file_path, "w") as f: - json.dump(train_obj, f, indent=4) + json.dump(project_gt, f, indent=4) def prepare_memory_replay( @@ -228,11 +309,10 @@ def prepare_memory_replay( superanimal_name: str, model_name: str, device: str, - max_individuals=3, - trainingsetindex: int = 0, - train_file="train.json", - pose_threshold=0.1, - customized_pose_checkpoint=None, + max_individuals: int = 3, + train_file: str = "train.json", + pose_threshold: float = 0.1, + customized_pose_checkpoint: str | None = None, ): """TODO: Documentation""" @@ -264,7 +344,10 @@ def prepare_memory_replay( memory_replay_folder = model_folder / "memory_replay" temp_dataset.materialize( - memory_replay_folder, framework="coco", append_image_id=False + memory_replay_folder, + framework="coco", + append_image_id=False, + no_image_copy=True, # use the images in the labeled-data folder ) original_model_config = af.read_config( @@ -290,19 +373,19 @@ def prepare_memory_replay( ) dataset = COCOPoseDataset(memory_replay_folder, "memory_replay_dataset") - conversion_table_path = dlc_proj_root / "memory_replay" / "conversion_table.csv" - # here we project the original DLC projects to superanimal space and save them into a coco project format + # here we project the original DLC projects to superanimal space and save them into + # a coco project format dataset.project_with_conversion_table(str(conversion_table_path)) - dataset.materialize(memory_replay_folder, deepcopy=False, framework="coco") - - # then in this function, we do pseudo label to match prediction and gts to create memory-replay dataset that will be named memory_replay_train.json - memory_replay_train_file = os.path.join( - memory_replay_folder, "annotations", "memory_replay_train.json" + dataset.materialize( + memory_replay_folder, framework="coco", deepcopy=False, no_image_copy=True, ) + # then in this function, we do pseudo label to match prediction and gts to create + # memory-replay dataset that will be named memory_replay_train.json prepare_memory_replay_dataset( + dlc_proj_root, memory_replay_folder, superanimal_name, model_name, diff --git a/deeplabcut/utils/pseudo_label.py b/deeplabcut/utils/pseudo_label.py index 5e279e7ed9..a44f922b40 100644 --- a/deeplabcut/utils/pseudo_label.py +++ b/deeplabcut/utils/pseudo_label.py @@ -145,30 +145,43 @@ def plot_cost_matrix( def keypoint_matching( - config_path, - superanimal_name, - model_name, - device=None, - train_file="train.json", - pose_threshold=0.1, + config_path: str | Path, + superanimal_name: str, + model_name: str, + copy_images: bool = False, + device: str | None = None, + train_file: str = "train.json", ): - - cfg = af.read_config(config_path) - - trainIndex = 0 - - dlc_proj_root = str(Path(config_path).parent) + """Runs the keypoint matching algorithm for a DeepLabCut project + + Matches project keypoints to SuperAnimal keypoints automatically, by running + SuperAnimal inference on all images in the dataset + + Args: + config_path: The path of the DeepLabCut project configuration file. + superanimal_name: SuperAnimal dataset with which to run keypoint matching. + model_name: SuperAnimal model with which to run keypoint matching + copy_images: When False, symlinks are created for the dataset used for keypoint + matching. Otherwise, images are copied from the `labeled-data` folder to the + folder used for keypoint matching. + device: The device on which to run keypoint matching. + train_file: The name of the file containing the labels to output. + """ + config_path = Path(config_path) + cfg = af.read_config(str(config_path)) + dlc_proj_root = config_path.parent if "individuals" in cfg: - temp_dataset = MaDLCDataFrame(dlc_proj_root, "temp_dataset") + temp_dataset = MaDLCDataFrame(str(dlc_proj_root), "temp_dataset") max_individuals = len(cfg["individuals"]) else: - temp_dataset = SingleDLCDataFrame(dlc_proj_root, "temp_dataset") + temp_dataset = SingleDLCDataFrame(str(dlc_proj_root), "temp_dataset") max_individuals = 1 - memory_replay_folder = Path(dlc_proj_root) / "memory_replay" - - temp_dataset.materialize(str(memory_replay_folder), framework="coco") + memory_replay_folder = dlc_proj_root / "memory_replay" + temp_dataset.materialize( + str(memory_replay_folder), framework="coco", deepcopy=copy_images + ) # inferencing the train set ( @@ -186,7 +199,6 @@ def keypoint_matching( individuals = [f"animal{i}" for i in range(max_individuals)] config["individuals"] = individuals - num_bodyparts = len(config["bodyparts"]) train_file_path = os.path.join(memory_replay_folder, "annotations", train_file) pose_runner, detector_runner = get_inference_runners( @@ -204,30 +216,29 @@ def keypoint_matching( images = train_obj["images"] annotations = train_obj["annotations"] categories = train_obj["categories"] - imagename2id = {} - imageid2name = {} - imagename2gt = defaultdict(list) + image_name_to_id = {} + image_id_to_name = {} + + image_name_to_gt = defaultdict(list) + image_name_to_bbox = defaultdict(list) + image_id_to_annotations = defaultdict(list) for image in images: # this only works with relative path as the testing image can be at a different folder - imagename = image["file_name"].split(os.sep)[-1] - imagename2id[imagename] = image["id"] - imageid2name[image["id"]] = imagename - - imagename2bbox = defaultdict(list) + name = image["file_name"].split(os.sep)[-1] + image_name_to_id[name] = image["id"] + image_id_to_name[image["id"]] = name for anno in annotations: - imagename = imageid2name[anno["image_id"]] - imagename2gt[imagename].append(anno) - imagename2bbox[imagename].append(anno["bbox"]) + name = image_id_to_name[anno["image_id"]] + image_name_to_gt[name].append(anno) + image_name_to_bbox[name].append(anno["bbox"]) - imageid2annotations = defaultdict(list) - - imageids = list(imagename2id.values()) - for annotation in annotations: - image_id = annotation["image_id"] - if annotation["image_id"] in imageids: - imageid2annotations[image_id].append(annotation) + image_ids = set(image_name_to_id.values()) + for anno in annotations: + image_id = anno["image_id"] + if anno["image_id"] in image_ids: + image_id_to_annotations[image_id].append(anno) # need to support more image types image_extensions = ["*.png", "*.jpg", "*.jpeg", "*.bmp", "*.gif", "*.tiff"] @@ -238,19 +249,15 @@ def keypoint_matching( ) corresponded_images = [] - for image in images_in_folder: image_path = image - imagename = image.split(os.sep)[-1] - if imagename in imagename2id: + name = image.split(os.sep)[-1] + if name in image_name_to_id: corresponded_images.append(image_path) images = corresponded_images - - bbox_predictions = detector_runner.inference(images=images) - bbox_gts = [ - {"bboxes": np.array(imagename2bbox[image.split(os.sep)[-1]])} + {"bboxes": np.array(image_name_to_bbox[image.split(os.sep)[-1]])} for image in images ] @@ -260,16 +267,14 @@ def keypoint_matching( predictions = pose_runner.inference(pose_inputs) with open(str(memory_replay_folder / "pseudo_predictions.json"), "w") as f: - json.dump(pose_inputs, f, cls=NumpyEncoder) assert len(images) == len(predictions) - imagename2prediction = {} - + image_name_to_pred = {} for image_path, prediction in zip(images, predictions): - imagename = image_path.split(os.sep)[-1] - imagename2prediction[imagename] = prediction + name = image_path.split(os.sep)[-1] + image_name_to_pred[name] = prediction pred_keypoint_names = config["bodyparts"] num_pred_keypoints = len(pred_keypoint_names) @@ -279,10 +284,10 @@ def keypoint_matching( match_matrix = np.zeros((num_pred_keypoints, num_gt_keypoints)) match_dict = defaultdict(lambda: defaultdict(int)) - for imagename, gts in imagename2gt.items(): + for name, gts in image_name_to_gt.items(): bbox_gts = [np.array(gt["bbox"]) for gt in gts] bbox_gts = [xywh2xyxy(e) for e in bbox_gts] - prediction = imagename2prediction[imagename] + prediction = image_name_to_pred[name] bbox_preds = [xywh2xyxy(pred) for pred in prediction["bboxes"]] optimal_pred_indices = optimal_match(bbox_gts, bbox_preds) @@ -293,16 +298,11 @@ def keypoint_matching( optimal_index = optimal_pred_indices[idx] matched_gt = np.array(gts[idx]["keypoints"]) matched_pred = prediction["bodyparts"][optimal_index] - bbox_gt = bbox_gts[idx] - bbox_pred = bbox_preds[idx] matched_gt = matched_gt.reshape(num_gt_keypoints, -1) matched_pred = matched_pred.reshape(num_pred_keypoints, -1) - gt_kpt_ids = np.arange(matched_gt.shape[0]) - pred_kpt_ids = np.arange(matched_pred.shape[0]) pair_distance = cdist(matched_pred, matched_gt) row_ind, column_ind = linear_sum_assignment(pair_distance) - original_gt_matched_indices = matched_gt[column_ind] for row, column in zip(row_ind, column_ind): pred_kpt_name = pred_keypoint_names[row] anno_kpt_name = gt_keypoint_names[column] diff --git a/examples/COLAB/COLAB_YOURDATA_SuperAnimal.ipynb b/examples/COLAB/COLAB_YOURDATA_SuperAnimal.ipynb new file mode 100644 index 0000000000..7d50b06351 --- /dev/null +++ b/examples/COLAB/COLAB_YOURDATA_SuperAnimal.ipynb @@ -0,0 +1,2164 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "view-in-github" + }, + "source": [ + "\"Open" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "5SSZpZUu0Z4S" + }, + "source": [ + "# DeepLabCut Model Zoo: SuperAnimal models\n", + "\n", + "![alt text](https://images.squarespace-cdn.com/content/v1/57f6d51c9f74566f55ecf271/1616492373700-PGOAC72IOB6AUE47VTJX/ke17ZwdGBToddI8pDm48kB8JrdUaZR-OSkKLqWQPp_YUqsxRUqqbr1mOJYKfIPR7LoDQ9mXPOjoJoqy81S2I8N_N4V1vUb5AoIIIbLZhVYwL8IeDg6_3B-BRuF4nNrNcQkVuAT7tdErd0wQFEGFSnBqyW03PFN2MN6T6ry5cmXqqA9xITfsbVGDrg_goIDasRCalqV8R3606BuxERAtDaQ/modelzoo.png?format=1000w)\n", + "\n", + "# 🦄 SuperAnimal in DeepLabCut PyTorch! 🔥\n", + "\n", + "This notebook demos how to use our SuperAnimal models within DeepLabCut 3.0! Please read more in [Ye et al. Nature Communications 2024](https://www.nature.com/articles/s41467-024-48792-2) about the available SuperAnimal models, and follow along below!\n", + "\n", + "### **Let's get going: install the latest version of DeepLabCut into COLAB:**\n", + "\n", + "*Also, be sure you are connected to a GPU: go to menu, click Runtime > Change Runtime Type > select \"GPU\"*\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/", + "height": 1000 + }, + "id": "AjET5cJE5UYM", + "jupyter": { + "outputs_hidden": true + }, + "outputId": "290a589f-a063-4933-d315-e13052ec1024" + }, + "outputs": [], + "source": [ + "!pip install \"git+https://github.com/DeepLabCut/DeepLabCut.git@pytorch_dlc#egg=deeplabcut[modelzoo]\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "5h0vq6E50Z4W" + }, + "source": [ + "**PLEASE, click \"restart runtime\" from the output above before proceeding!**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "LvnlIvQm0Z4X", + "jupyter": { + "outputs_hidden": true + }, + "outputId": "ef4fd2ed-4569-41d4-b78a-8bf5ae9a0e6b" + }, + "outputs": [], + "source": [ + "import os\n", + "from pathlib import Path\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import pandas as pd\n", + "from PIL import Image\n", + "\n", + "import deeplabcut\n", + "from deeplabcut.pose_estimation_pytorch.apis import (\n", + " superanimal_analyze_images,\n", + ")\n", + "from deeplabcut.core.weight_init import WeightInitialization\n", + "from deeplabcut.core.engine import Engine\n", + "from deeplabcut.modelzoo.utils import (\n", + " create_conversion_table,\n", + " read_conversion_table_from_csv,\n", + ")\n", + "from deeplabcut.modelzoo.video_inference import video_inference_superanimal\n", + "from deeplabcut.utils.pseudo_label import keypoint_matching" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "UeXjmtu40Z4X" + }, + "source": [ + "## Zero-shot Image & Video Inference\n", + "SuperAnimal models are foundation animal pose models. They can be used for zero-shot predictions without further training on the data.\n", + "In this section, we show how to use SuperAnimal models to predict pose from images (given an image folder) and output the predicted images (with pose) into another destination folder." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "FvFzntDMxPoL" + }, + "source": [ + "### Zero-shot image inference\n", + "\n", + "If you have a single Image you want to test, upload it here!" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "NbDsZQfsxPoL" + }, + "source": [ + "#### Upload the images you want to predict" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "c4yfTj7r0Z4Y", + "jupyter": { + "outputs_hidden": true + } + }, + "outputs": [], + "source": [ + "from google.colab import files\n", + "\n", + "uploaded = files.upload()\n", + "for filepath, content in uploaded.items():\n", + " print(f\"User uploaded file '{filepath}' with length {len(content)} bytes\")\n", + "image_path = os.path.abspath(filepath)\n", + "image_name = os.path.splitext(image_path)[0]\n", + "\n", + "# If this cell fails (e.g., when using Safari in place of Google Chrome),\n", + "# manually upload your video via the Files menu to the left\n", + "# and define `video_path` yourself with right click > copy path on the video." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Jashzdjb0Z4Y" + }, + "source": [ + "#### Select a SuperAnimal name and corresponding model architecture\n", + "\n", + "Check Our Docs on [SuperAnimals](https://github.com/DeepLabCut/DeepLabCut/blob/main/docs/ModelZoo.md) to learn more!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "uH9LXig90Z4Y", + "jupyter": { + "outputs_hidden": true + } + }, + "outputs": [], + "source": [ + "# @markdown ---\n", + "# @markdown SuperAnimal Configurations\n", + "superanimal_name = \"superanimal_topviewmouse\" #@param [\"superanimal_topviewmouse\", \"superanimal_quadruped\"]\n", + "model_name = \"hrnetw32\" #@param [\"hrnetw32\"]\n", + "\n", + "# @markdown ---\n", + "# @markdown What is the maximum number of animals you expect to have in an image\n", + "max_individuals = 3 # @param {type:\"slider\", min:1, max:30, step:1}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "OmJtVmHq0Z4Y", + "jupyter": { + "outputs_hidden": true + } + }, + "outputs": [], + "source": [ + "# Note you need to enter max_individuals correctly to get the correct number of predictions in the image.\n", + "_ = superanimal_analyze_images(\n", + " superanimal_name,\n", + " model_name,\n", + " image_path,\n", + " max_individuals,\n", + " out_folder=\"/content/\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "6VEjHu-00Z4Y" + }, + "source": [ + "### Zero-shot Video Inference\n", + "\n", + "This can be done with or without video adaptation (faster, but not self-supervised fine-tuned on your data!)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "qGoAhxZOxPoM" + }, + "source": [ + "#### Upload a video you want to predict" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "PK3efA0I0Z4Y", + "jupyter": { + "outputs_hidden": true + } + }, + "outputs": [], + "source": [ + "from google.colab import files\n", + "\n", + "uploaded = files.upload()\n", + "for filepath, content in uploaded.items():\n", + " print(f\"User uploaded file '{filepath}' with length {len(content)} bytes\")\n", + "video_path = os.path.abspath(filepath)\n", + "video_name = os.path.splitext(video_path)[0]\n", + "\n", + "# If this cell fails (e.g., when using Safari in place of Google Chrome),\n", + "# manually upload your video via the Files menu to the left\n", + "# and define `video_path` yourself with right click > copy path on the video." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "JoA-RATSICj_" + }, + "source": [ + "#### Choose the superanimal and the model name" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "OiRAP9XD0Z4Z", + "jupyter": { + "outputs_hidden": true + } + }, + "outputs": [], + "source": [ + "# @markdown ---\n", + "# @markdown SuperAnimal Configurations\n", + "superanimal_name = \"superanimal_topviewmouse\" #@param [\"superanimal_topviewmouse\", \"superanimal_quadruped\"]\n", + "model_name = \"hrnetw32\" #@param [\"hrnetw32\"]\n", + "\n", + "# @markdown ---\n", + "# @markdown What is the maximum number of animals you expect to have in an image\n", + "max_individuals = 3 # @param {type:\"slider\", min:1, max:30, step:1}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Zv3v0QgSJNOg" + }, + "source": [ + "#### Zero-shot Video Inference without video adaptation\n", + "\n", + "The labeled video (and pose predictions for the video) are saved in `\"/content/\"`, with the labeled video name being `{your_video_name}_superanimal_{superanimal_name}_hrnetw32_labeled.mp4`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "poqynL0UJTBp", + "jupyter": { + "outputs_hidden": true + } + }, + "outputs": [], + "source": [ + "_ = video_inference_superanimal(\n", + " videos=video_path,\n", + " superanimal_name=f\"{superanimal_name}_{model_name}\",\n", + " video_adapt=False,\n", + " max_individuals=max_individuals,\n", + " dest_folder=\"/content/\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Z8Z5GSti0Z4Z" + }, + "source": [ + "#### Zero-shot Video Inference with video adaptation (unsupervised)\n", + "\n", + "The labeled video (and pose predictions for the video) are saved in `\"/content/\"`, with the labeled video name being `{your_video_name}_superanimal_{superanimal_name}_hrnetw32_labeled_after_adapt.mp4`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "5mhOmtzw0Z4Z", + "jupyter": { + "outputs_hidden": true + } + }, + "outputs": [], + "source": [ + "_ = video_inference_superanimal(\n", + " videos=[video_path],\n", + " superanimal_name=f\"{superanimal_name}_{model_name}\",\n", + " video_adapt=True,\n", + " max_individuals=max_individuals,\n", + " pseudo_threshold=0.1,\n", + " bbox_threshold=0.9,\n", + " detector_epochs=1,\n", + " pose_epochs=1,\n", + " dest_folder=\"/content/\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "br3pwGf40Z4a" + }, + "source": [ + "## Training with SuperAnimal\n", + "\n", + "In this section, we compare different ways to train models in DeepLabCut 3.0, with or without using SuperAnimal-pretrained models.\n", + "You can compare the evaluation results and get a sense of each baseline. We have following baselines:\n", + "\n", + "- ImageNet transfer learning (training without superanimal)\n", + "- SuperAnimal transfer learning (baseline 1)\n", + "- SuperAnimal naive fine-tuning (baseline 2)\n", + "- SuperAnimal memory-replay fine-tuning (baseline3)\n", + "\n", + "This is done on one of your DeepLabCut projects! If you don't have a DeepLabCut project that you can use SuperAnimal models with, you can always using the example openfield dataset [available in the DeepLabCut repository](https://github.com/DeepLabCut/DeepLabCut/tree/main/examples/openfield-Pranav-2018-10-30) or the Tri-Mouse dataset available on [Zenodo](https://zenodo.org/records/5851157)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "yPy5VgDDhD6o" + }, + "source": [ + "### Preparing the DeepLabCut Project\n", + "\n", + "First, place your DeepLabCut project folder into you google drive! \"i.e. move the folder named \"Project-YourName-TheDate\" into Google Drive." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "SXzBBV8ehDR9", + "outputId": "90d61c19-400b-4e5d-8ac9-63680d72cdb5" + }, + "outputs": [], + "source": [ + "# Now, let's link to your GoogleDrive. Run this cell and follow the\n", + "# authorization instructions:\n", + "\n", + "from google.colab import drive\n", + "drive.mount('/content/drive')" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "-QmTftBMo4h6" + }, + "source": [ + "You will need to edit the project path in the config.yaml file to be set to your Google Drive link!\n", + "\n", + "Typically, this will be in the format: `/content/drive/MyDrive/yourProjectFolderName`. You can obtain this path by going to the file navigator in the left pane, finding your DeepLabCut project folder, clicking on the vertical `...` next to the folder name and selecting \"Copy path\".\n", + "\n", + "If the `drive` folder is not immediately visible after mounting the drive, refresh the available files!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "_iFFEYAB7Uum" + }, + "outputs": [], + "source": [ + "# TODO: Update the `project_path` to be the path of your DeepLabCut project!\n", + "project_path = Path(\"/content/drive/MyDrive/my-project-2024-07-17\")\n", + "config_path = str(project_path / \"config.yaml\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "HZTG3Eo475w0" + }, + "source": [ + "Then, use the panel below to select the appropriate SuperAnimal model for your project (don't forget to run the cell)!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "t8NtCy1Jo0bu" + }, + "outputs": [], + "source": [ + "# @markdown ---\n", + "# @markdown SuperAnimal Configurations\n", + "superanimal_name = \"superanimal_topviewmouse\" #@param [\"superanimal_topviewmouse\", \"superanimal_quadruped\"]\n", + "model_name = \"hrnetw32\" #@param [\"hrnetw32\"]" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "BPvoL9uZ0Z4a" + }, + "source": [ + "### Comparison between different training baselines\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "eVmpaLdB0Z4a" + }, + "source": [ + "Definition of data split: the unique combination of training images and testing images.\n", + "We create a data split named split 0. All baselines will share the data split to make fair comparisons.\n", + "- split 0 -> shared by all baselines\n", + "- shuffle 0 (split0) -> imagenet transfer learning\n", + "- shuffle 1 (split0) -> superanimal transfer learning\n", + "- shuffle 2 (split0) -> superanimal naive fine-tuning\n", + "- shuffle 3 (split0) -> superanimal memory-replay fine-tuning" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "WofR2jytxPoR" + }, + "source": [ + "### What is the difference between baselines?\n", + "\n", + "**Transfer learning** For canonical task-agnostic transfer learning,\n", + "the encoder learns universal visual features from a large pre-training dataset, and a randomly\n", + "initialized decoder is used to learn the pose from the downstream dataset.\n", + "\n", + "**Fine-tuning** For task aware\n", + "fine-tuning, both encoder and decoder learn task-related visual-pose features\n", + "in the pre-training datasets, and the decoder is fine-tuned to update pose\n", + "priors in downstream datasets. Crucially, the network has pose-estimation-specific\n", + "weights\n", + "\n", + "**ImageNet transfer-learning** The encoder was pre-trained from ImageNet. The decoder is trained from scratch in the downstream tasks\n", + "\n", + "**SuperAnimal transfer-learning** The encoder was pre-trained first from ImageNet, then in pose datasets we colleceted. Then decoder is trained from scratch in downstream tasks.\n", + "\n", + "**SuperAnimal naive fine-tuning** Both the encoder and the decoder were pre-trained in pose datasets we collected. In downstream datsets, we only finetune convolutional channels that correspond to the annotated keypoints in the downstream datasets. This introduces catastrophic forgetting in keypoints that are not annotated in the downstream datasets.\n", + "\n", + "**SuperAnimal memory-replay fine-tuning** If we apply fine-tuning with SuperAnimal without further cares, the models will forget about keypoints that are not annotated in the downstream datasets. To mitigate this, we mix the annotations and zero-shot predictions of SuperAnimal models to create a dataset that 'replays' the memory of the SuperAnimal keypoints.\n", + "\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "AgIsUu6v0Z4a", + "jupyter": { + "outputs_hidden": true + } + }, + "outputs": [], + "source": [ + "imagenet_transfer_learning_shuffle = 0\n", + "superanimal_transfer_learning_shuffle = 1\n", + "superanimal_naive_finetune_shuffle = 2\n", + "superanimal_memory_replay_shuffle = 3" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "kuKcxM8F0Z4a", + "jupyter": { + "outputs_hidden": true + }, + "outputId": "c7df2943-1e2c-4b85-c20d-8b94a8aabd75" + }, + "outputs": [], + "source": [ + "deeplabcut.create_training_dataset(\n", + " config_path,\n", + " Shuffles=[imagenet_transfer_learning_shuffle],\n", + " net_type=\"top_down_hrnet_w32\",\n", + " engine=Engine.PYTORCH,\n", + " userfeedback=False,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "_6RncQbr0Z4a" + }, + "source": [ + "### ImageNet transfer learning\n", + "\n", + "Historically, the transfer learning using ImageNet weights strategies assumed no “animal pose task priors” in the pretrained\n", + "model, a paradigm adopted from previous task-agnostic transfer learning.\n", + "\n", + "You can change the number of epochs you want to train for. How long training will take depends on many parameters, including the number of images in your dataset, the resolution of the images, and the number of epochs you train for." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/", + "height": 1000, + "referenced_widgets": [ + "7ed11ae2a4be462da84ff716e0725af0", + "0f0ed94a863f49b9b85d0a18fa8ce2a5", + "343f2670d37c4bf18859238c3d81d419", + "d104ae21091e4f10a7de18e191b9f04d", + "5dcbd8f3fb6148cca6cfc72b20ce49bd", + "e1675e53ca9a4da8acf6c16fba7a2578", + "3d2996e10f96404baf24d2c4215b75a1", + "b988f87e676840ee98daa3d996c9ddbc", + "1779b84e748b4989a8ed53434c30016f", + "d37cf6fe7c444bc2a2568c3407389ea8", + "2cef5e028d2e40a6bba7400be922d0c2" + ] + }, + "id": "H2z8kM340Z4a", + "jupyter": { + "outputs_hidden": true + }, + "outputId": "75cc2c95-2ac7-4354-9134-4847937e15ce" + }, + "outputs": [], + "source": [ + "# Note we skip the detector training to save time.\n", + "# For Top-Down models, the evaluation is by default using ground-truth bounding\n", + "# boxes. But to train a model that can be used to inference videos and images,\n", + "# you have to set detector_epochs > 0.\n", + "\n", + "deeplabcut.train_network(\n", + " config_path,\n", + " detector_epochs=0,\n", + " epochs=50,\n", + " save_epochs=10,\n", + " batch_size=64, # if you get a CUDA OOM error when training on a GPU, reduce to 32, 16, ...!\n", + " display_iters=10,\n", + " shuffle=imagenet_transfer_learning_shuffle,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "J-udMck7nDbG" + }, + "source": [ + "Now let's evaluate the performance of our trained models." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "TDHMdKz4m_16", + "jupyter": { + "outputs_hidden": true + }, + "outputId": "1d38fb84-7f4c-45d1-dbcd-fd7117ca4dad" + }, + "outputs": [], + "source": [ + "deeplabcut.evaluate_network(config_path, Shuffles=[imagenet_transfer_learning_shuffle])" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "0GIFWU-MxPoR" + }, + "source": [ + "### Transfer learning with SuperAnimal weights\n", + "\n", + "First, we prepare training shuffle for transfer-learning with SuperAnimal weights. As we've already create a shuffle with a train/test split that we want to reuse, we use `deeplabcut.create_training_dataset_from_existing_split` to keep the same train/test indices as in the ImageNet transfer learning shuffle.\n", + "\n", + "We specify that we want to initialize the model weights with the selected SuperAnimal model, but without keeping the decoding layers (this is called transfer learning)!\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "wOSdZQtOp8qa", + "jupyter": { + "outputs_hidden": true + }, + "outputId": "ea721606-ea9f-444b-cdae-f62cf0ad30be" + }, + "outputs": [], + "source": [ + "weight_init = WeightInitialization(\n", + " dataset=superanimal_name,\n", + " with_decoder=False,\n", + ")\n", + "\n", + "deeplabcut.create_training_dataset_from_existing_split(\n", + " config_path,\n", + " from_shuffle=imagenet_transfer_learning_shuffle,\n", + " shuffles=[superanimal_transfer_learning_shuffle],\n", + " engine=Engine.PYTORCH,\n", + " net_type=\"top_down_hrnet_w32\",\n", + " weight_init=weight_init,\n", + " userfeedback=False,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "3qFxlRHixPoR" + }, + "source": [ + "Then, we launch the training for transfer-learning with SuperAnimal weights." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/", + "height": 1000, + "referenced_widgets": [ + "9a996c8dc3b34bc5b8805b3687e22b27", + "d012b421c189412dabeac84cba4164a7", + "1abff22a7c9a416d9166e6b150612171", + "7271412c1f0141649a7300dbce2b003c", + "3c011813d7cb48588a8d236785d9c24f", + "3ea385fe815f4e50a0b81ec299040314", + "fe59f6c5ed7b4e2cb87bb60224acdaba", + "04370d8302c04c5ca6a351383126193f", + "d67c4871543e405fbb576a55f8c9048a", + "a6cb25fa67ef4733a720960b3fc8213c", + "b73b1b64620d492dbc4eaf4bd83ca23a", + "dccbe277cc084ed6aa0b329067b5c69c", + "c8b57833d3f946abae69b84075345a54", + "bee292213d8645618536fcdf6a491d83", + "fbbc8c5b20c7423fb21b74296e0eeb28", + "ff0c737c49624b1ea27588611951fc84", + "42874cdab4be4dc38b0c33775b27d98c", + "e3a185abf8a04edabf32d58bdee10dd1", + "7cdcbbf9cb694dbf949e8b7eea8e7836", + "2ec06260b237411cabd3de7c37e03b1b", + "9f8009429aa34b40a65c998230f20c99", + "2a3abfe7867641db9fbfe3ee76854bf4" + ] + }, + "id": "W60UgRQWqghn", + "jupyter": { + "outputs_hidden": true + }, + "outputId": "18b931b8-98f4-4539-bf82-1910ff5b7f70" + }, + "outputs": [], + "source": [ + "deeplabcut.train_network(\n", + " config_path,\n", + " detector_epochs=0,\n", + " epochs=50,\n", + " save_epochs=10,\n", + " batch_size=64, # if you get a CUDA OOM error when training on a GPU, reduce to 32, 16, ...!\n", + " display_iters=10,\n", + " shuffle=superanimal_transfer_learning_shuffle,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "XzOWKiOixPoR" + }, + "source": [ + "Finally, we evaluate the model obtained by transfer-learning with SuperAnimal weights." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "jpO3aIAIsWbz", + "jupyter": { + "outputs_hidden": true + }, + "outputId": "30415e5b-8011-4651-af77-a781ea2b5af7" + }, + "outputs": [], + "source": [ + "deeplabcut.evaluate_network(config_path, Shuffles=[superanimal_transfer_learning_shuffle])" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "_Es6RR-_0Z4b" + }, + "source": [ + "### Fine-tuning with SuperAnimal (without keeping full SuperAnimal keypoints)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "6oo9oJ8XyZrn" + }, + "source": [ + "#### Setup the weight init and dataset\n", + "\n", + "First we do keypoint matching. This steps make it possible to understand the correspondance between the existing annotations and SuperAnimal annotations. This step produces 3 outputs\n", + "- The confusion matrix\n", + "- The conversion table\n", + "- Pseudo predictions over the whole dataset" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "fRm62Ji_xPoS" + }, + "source": [ + "#### What is keypoint matching?\n", + "\n", + "Because SuperAnimal models have their pre-defined keypoints that are potentially different from your annotations, we porposed this algorithm to minimize the gap between the model and the dataset. We use our model to perform zero-shot inference on the whole dataset. This gives pairs of predictions and ground truth for every image. Then, we cast the matching between models’ predictions (2D coordinates)\n", + "and ground truth as bipartitematching using the Euclidean distance as the cost between paired of keypoints. We then solve the matching using the Hungarian algorithm. Thus for every image, we end up getting a matching matrix where 1 counts formatch and 0 counts for non-matching. Because the models’ predictions can be noisy from image to image, we average the aforementioned matching matrix across all the images and perform another bipartite matching, resulting in the final keypoint conversion table between the model and the dataset. Note that the quality of thematching will impact the performance\n", + "of the model, especially for zero-shot. In the case where, e.g., the annotation nose is mistakenly converted to keypoint tail and vice versa, the model will have to unlearn the channel that corresponds to nose and tail (see also case study in Mathis et al.)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/", + "height": 1000 + }, + "id": "vEHeuKSKyjA6", + "jupyter": { + "outputs_hidden": true + }, + "outputId": "5863a81e-e0b9-48c7-f2f9-de14d38e805e" + }, + "outputs": [], + "source": [ + "keypoint_matching(\n", + " config_path,\n", + " superanimal_name,\n", + " model_name,\n", + " copy_images=True,\n", + ")\n", + "\n", + "conversion_table_path = project_path / \"memory_replay\" / \"conversion_table.csv\"\n", + "confusion_matrix_path = project_path / \"memory_replay\" / \"confusion_matrix.png\"\n", + "\n", + "# You can visualize the pseudo predictions, or do pose embedding clustering etc.\n", + "pseudo_prediction_path = project_path / \"memory_replay\" / \"pseudo_predictions.json\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "sA8yyLgs0zoO" + }, + "source": [ + "#### Display the confusion matrix\n", + "\n", + "The x axis lists the keypoints in the existing annotations. The y axis lists the keypoints in SuperAnimal keypoint space. Darker color encodes stronger correspondance between the human annotation and SuperAnimal annotations." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "luDxpD9H0zYZ", + "jupyter": { + "outputs_hidden": true + } + }, + "outputs": [], + "source": [ + "confusion_matrix_image = Image.open(confusion_matrix_path)\n", + "\n", + "plt.imshow(confusion_matrix_image)\n", + "plt.axis('off') # Hide the axes for better view\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "i0QWikYmy_Mj" + }, + "source": [ + "#### Display the conversion table\n", + "The gt columns represents the keypoint names in the existing dataset. The MasterName represents the correspoinding keypoints in SuperAnimal keypoint space." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "CeA-NzDMynYV", + "jupyter": { + "outputs_hidden": true + } + }, + "outputs": [], + "source": [ + "df = pd.read_csv(conversion_table_path)\n", + "df = df.dropna()\n", + "\n", + "df" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "GkfIo8zTxPoS" + }, + "source": [ + "#### Prepare the training shuffle and weight initialization for (naive) fine-tuning with SuperAnimal weights" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "xEeM_hrOu6k8", + "jupyter": { + "outputs_hidden": true + } + }, + "outputs": [], + "source": [ + "table = create_conversion_table(\n", + " config=config_path,\n", + " super_animal=superanimal_name,\n", + " project_to_super_animal=read_conversion_table_from_csv(conversion_table_path),\n", + ")\n", + "\n", + "weight_init = WeightInitialization(\n", + " dataset=superanimal_name,\n", + " with_decoder=True,\n", + " conversion_array=table.to_array()\n", + ")\n", + "\n", + "deeplabcut.create_training_dataset_from_existing_split(\n", + " config_path,\n", + " from_shuffle=imagenet_transfer_learning_shuffle,\n", + " shuffles=[superanimal_naive_finetune_shuffle],\n", + " engine=Engine.PYTORCH,\n", + " net_type=\"top_down_hrnet_w32\",\n", + " weight_init=weight_init,\n", + " userfeedback=False,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "gZx6nr-ExPoS" + }, + "source": [ + "#### Launch the training for (naive) fine-tuning with SuperAnimal" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "c3XAr6uRyXOD", + "jupyter": { + "outputs_hidden": true + } + }, + "outputs": [], + "source": [ + "deeplabcut.train_network(\n", + " config_path,\n", + " detector_epochs=0,\n", + " epochs=50,\n", + " save_epochs=10,\n", + " batch_size=64, # if you get a CUDA OOM error when training on a GPU, reduce to 32, 16, ...!\n", + " display_iters=10,\n", + " shuffle=superanimal_naive_finetune_shuffle,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "oXuRshzhxPoS" + }, + "source": [ + "#### Evaluate the model obtained by (naive) fine-tuning with SuperAnimal" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "VXfdKS-H2yqw", + "jupyter": { + "outputs_hidden": true + } + }, + "outputs": [], + "source": [ + "deeplabcut.evaluate_network(\n", + " config_path,\n", + " Shuffles=[superanimal_naive_finetune_shuffle],\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "_nUAMlbZ0Z4b" + }, + "source": [ + "### Memory-replay fine-tuning with SuperAnimal (keeping full SuperAnimal keypoints)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "n6HPu6RaxPoS" + }, + "source": [ + "**Catastrophic forgetting** describes a\n", + "classic problemin continual learning. Indeed, amodel gradually loses\n", + "its ability to solve previous tasks after it learns to solve new ones.\n", + "Fine-tuning a SuperAnimal models falls into the category of continual\n", + "learning: the downstream dataset defines potentially different\n", + "keypoints than those learned by the models. Thus, the models might\n", + "forget the keypoints they learned and only pick up those defined in the\n", + "target dataset. Here, retraining with the original dataset and the new\n", + "one, is not a feasible option as datasets cannot be easily shared and\n", + "more computational resources would be required.\n", + "To counter that, we treat zero-shot inference of the model as a\n", + "memory buffer that stores knowledge from the original model. When\n", + "we fine-tune a SuperAnimal model, we replace the model predicted\n", + "keypoints with the ground-truth annotations, resulting in hybrid\n", + "learning of old and new knowledge. The quality of the zero-shot predictions\n", + "can vary and we use the confidence of prediction (0.7) as a\n", + "threshold to filter out low-confidence predictions. With the threshold\n", + "set to 1, memory replay fine-tuning becomes naive-fine-tuning." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "CSLmjlCIxPoS" + }, + "source": [ + "#### Prepare training shuffle and weight initialization for memory-replay finetuning with SuperAnimal" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "BKEF76AI0Z4c", + "jupyter": { + "outputs_hidden": true + } + }, + "outputs": [], + "source": [ + "weight_init = WeightInitialization(\n", + " dataset=superanimal_name,\n", + " conversion_array=table.to_array(),\n", + " with_decoder=True,\n", + " memory_replay=True,\n", + ")\n", + "\n", + "deeplabcut.create_training_dataset_from_existing_split(\n", + " config_path,\n", + " from_shuffle=imagenet_transfer_learning_shuffle,\n", + " shuffles=[superanimal_memory_replay_shuffle],\n", + " engine=Engine.PYTORCH,\n", + " net_type=\"top_down_hrnet_w32\",\n", + " weight_init=weight_init,\n", + " userfeedback=False,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "MKwJiIyKxPoT" + }, + "source": [ + "#### Launch the training for memory-replay fine-tuning with SuperAnimal" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Ru8tIFmD2Mkv", + "jupyter": { + "outputs_hidden": true + } + }, + "outputs": [], + "source": [ + "deeplabcut.train_network(\n", + " config_path,\n", + " detector_epochs=0,\n", + " epochs=50,\n", + " save_epochs=10,\n", + " batch_size=64, # if you get a CUDA OOM error when training on a GPU, reduce to 32, 16, ...!\n", + " display_iters=10,\n", + " shuffle=superanimal_memory_replay_shuffle,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "i-2MBRDjxPoT" + }, + "source": [ + "#### Evaluate the model obtained by memory-replay finetuning with SuperAnimal" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "sfMcK3gq8WxZ", + "jupyter": { + "outputs_hidden": true + } + }, + "outputs": [], + "source": [ + "deeplabcut.evaluate_network(config_path, Shuffles=[superanimal_memory_replay_shuffle])" + ] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "collapsed_sections": [ + "UeXjmtu40Z4X", + "FvFzntDMxPoL", + "6VEjHu-00Z4Y" + ], + "gpuType": "T4", + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.13" + }, + "widgets": { + "application/vnd.jupyter.widget-state+json": { + "04370d8302c04c5ca6a351383126193f": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "0f0ed94a863f49b9b85d0a18fa8ce2a5": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_e1675e53ca9a4da8acf6c16fba7a2578", + "placeholder": "​", + "style": "IPY_MODEL_3d2996e10f96404baf24d2c4215b75a1", + "value": "model.safetensors: 100%" + } + }, + "1779b84e748b4989a8ed53434c30016f": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "ProgressStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "ProgressStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "bar_color": null, + "description_width": "" + } + }, + "1abff22a7c9a416d9166e6b150612171": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "FloatProgressModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "FloatProgressModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "ProgressView", + "bar_style": "success", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_04370d8302c04c5ca6a351383126193f", + "max": 159594859, + "min": 0, + "orientation": "horizontal", + "style": "IPY_MODEL_d67c4871543e405fbb576a55f8c9048a", + "value": 159594859 + } + }, + "2a3abfe7867641db9fbfe3ee76854bf4": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "2cef5e028d2e40a6bba7400be922d0c2": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "2ec06260b237411cabd3de7c37e03b1b": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "ProgressStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "ProgressStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "bar_color": null, + "description_width": "" + } + }, + "343f2670d37c4bf18859238c3d81d419": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "FloatProgressModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "FloatProgressModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "ProgressView", + "bar_style": "success", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_b988f87e676840ee98daa3d996c9ddbc", + "max": 165432914, + "min": 0, + "orientation": "horizontal", + "style": "IPY_MODEL_1779b84e748b4989a8ed53434c30016f", + "value": 165432914 + } + }, + "3c011813d7cb48588a8d236785d9c24f": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "3d2996e10f96404baf24d2c4215b75a1": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "3ea385fe815f4e50a0b81ec299040314": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "42874cdab4be4dc38b0c33775b27d98c": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "5dcbd8f3fb6148cca6cfc72b20ce49bd": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "7271412c1f0141649a7300dbce2b003c": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_a6cb25fa67ef4733a720960b3fc8213c", + "placeholder": "​", + "style": "IPY_MODEL_b73b1b64620d492dbc4eaf4bd83ca23a", + "value": " 160M/160M [00:00<00:00, 201MB/s]" + } + }, + "7cdcbbf9cb694dbf949e8b7eea8e7836": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "7ed11ae2a4be462da84ff716e0725af0": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HBoxModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HBoxModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HBoxView", + "box_style": "", + "children": [ + "IPY_MODEL_0f0ed94a863f49b9b85d0a18fa8ce2a5", + "IPY_MODEL_343f2670d37c4bf18859238c3d81d419", + "IPY_MODEL_d104ae21091e4f10a7de18e191b9f04d" + ], + "layout": "IPY_MODEL_5dcbd8f3fb6148cca6cfc72b20ce49bd" + } + }, + "9a996c8dc3b34bc5b8805b3687e22b27": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HBoxModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HBoxModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HBoxView", + "box_style": "", + "children": [ + "IPY_MODEL_d012b421c189412dabeac84cba4164a7", + "IPY_MODEL_1abff22a7c9a416d9166e6b150612171", + "IPY_MODEL_7271412c1f0141649a7300dbce2b003c" + ], + "layout": "IPY_MODEL_3c011813d7cb48588a8d236785d9c24f" + } + }, + "9f8009429aa34b40a65c998230f20c99": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "a6cb25fa67ef4733a720960b3fc8213c": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "b73b1b64620d492dbc4eaf4bd83ca23a": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "b988f87e676840ee98daa3d996c9ddbc": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "bee292213d8645618536fcdf6a491d83": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "FloatProgressModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "FloatProgressModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "ProgressView", + "bar_style": "success", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_7cdcbbf9cb694dbf949e8b7eea8e7836", + "max": 517816013, + "min": 0, + "orientation": "horizontal", + "style": "IPY_MODEL_2ec06260b237411cabd3de7c37e03b1b", + "value": 517816013 + } + }, + "c8b57833d3f946abae69b84075345a54": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_42874cdab4be4dc38b0c33775b27d98c", + "placeholder": "​", + "style": "IPY_MODEL_e3a185abf8a04edabf32d58bdee10dd1", + "value": "detector.pt: 100%" + } + }, + "d012b421c189412dabeac84cba4164a7": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_3ea385fe815f4e50a0b81ec299040314", + "placeholder": "​", + "style": "IPY_MODEL_fe59f6c5ed7b4e2cb87bb60224acdaba", + "value": "pose_model.pth: 100%" + } + }, + "d104ae21091e4f10a7de18e191b9f04d": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_d37cf6fe7c444bc2a2568c3407389ea8", + "placeholder": "​", + "style": "IPY_MODEL_2cef5e028d2e40a6bba7400be922d0c2", + "value": " 165M/165M [00:04<00:00, 41.1MB/s]" + } + }, + "d37cf6fe7c444bc2a2568c3407389ea8": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "d67c4871543e405fbb576a55f8c9048a": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "ProgressStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "ProgressStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "bar_color": null, + "description_width": "" + } + }, + "dccbe277cc084ed6aa0b329067b5c69c": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HBoxModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HBoxModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HBoxView", + "box_style": "", + "children": [ + "IPY_MODEL_c8b57833d3f946abae69b84075345a54", + "IPY_MODEL_bee292213d8645618536fcdf6a491d83", + "IPY_MODEL_fbbc8c5b20c7423fb21b74296e0eeb28" + ], + "layout": "IPY_MODEL_ff0c737c49624b1ea27588611951fc84" + } + }, + "e1675e53ca9a4da8acf6c16fba7a2578": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "e3a185abf8a04edabf32d58bdee10dd1": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "fbbc8c5b20c7423fb21b74296e0eeb28": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_9f8009429aa34b40a65c998230f20c99", + "placeholder": "​", + "style": "IPY_MODEL_2a3abfe7867641db9fbfe3ee76854bf4", + "value": " 518M/518M [00:05<00:00, 101MB/s]" + } + }, + "fe59f6c5ed7b4e2cb87bb60224acdaba": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "ff0c737c49624b1ea27588611951fc84": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + } + } + } + }, + "nbformat": 4, + "nbformat_minor": 1 +}