diff --git a/deeplabcut/pose_estimation_pytorch/README.md b/deeplabcut/pose_estimation_pytorch/README.md index dc1cd27bf3..09ee4113d4 100644 --- a/deeplabcut/pose_estimation_pytorch/README.md +++ b/deeplabcut/pose_estimation_pytorch/README.md @@ -321,9 +321,13 @@ project_root = Path("/path/to/my/COCOProject") train_json_filename = "train.json" test_json_filename = "test.json" -# Parse information about the project +# Parse information about the project. (Pass test_dict if you have a test set) train_dict = dlc_torch.COCOLoader.load_json(project_root, filename=train_json_filename) -max_num_individuals, bodyparts = dlc_torch.COCOLoader.get_project_parameters(train_dict) +test_dict = dlc_torch.COCOLoader.load_json(project_root, filename=test_json_filename) +max_num_individuals, bodyparts = dlc_torch.COCOLoader.get_project_parameters( + train_dict, + test_dict, +) # Generate a configuration file for your PyTorch model # In this case, it's for a Top-Down HRNet_w32 diff --git a/deeplabcut/pose_estimation_pytorch/data/cocoloader.py b/deeplabcut/pose_estimation_pytorch/data/cocoloader.py index 2f7ed00f36..f6aece9869 100644 --- a/deeplabcut/pose_estimation_pytorch/data/cocoloader.py +++ b/deeplabcut/pose_estimation_pytorch/data/cocoloader.py @@ -74,6 +74,49 @@ def __init__( if self.test_json_filename: self.test_json = self.load_json(self.project_root, self.test_json_filename) + self._validate_against_model_cfg() + + def _validate_against_model_cfg(self) -> None: + """Checks that the COCO annotations are compatible with `model_cfg.metadata`. + `model_cfg.metadata` is authoritative for `bodyparts` and `individuals` + (the pose model is built with these parameters, and the COCO JSON should match). + + Raises: + ValueError: If an image has more individuals than `model_cfg` supports, or + if the annotated bodyparts don't match `model_cfg.metadata.bodyparts`. + """ + meta = self.model_cfg.metadata + bodyparts = list(meta.bodyparts) + + for name, coco_json in (("train", self.train_json), ("test", self.test_json)): + if coco_json is None: + continue + + json_bodyparts = list(coco_json["categories"][0]["keypoints"]) + if json_bodyparts != bodyparts: + raise ValueError( + f"The bodyparts in {self.train_json_filename if name == 'train' else self.test_json_filename} " + f"({json_bodyparts}) don't match model_cfg.metadata.bodyparts ({bodyparts}). The order must " + "match exactly, as it determines which keypoint index is associated with which bodypart name." + ) + + observed = self._max_individuals_in_json(coco_json) + if observed > meta.num_individuals: + raise ValueError( + f"{self.train_json_filename if name == 'train' else self.test_json_filename} has an image " + f"with {observed} individuals, but model_cfg only supports {meta.num_individuals} " + f"(metadata.individuals={list(meta.individuals)}). Rebuild the model config with " + f"max_individuals >= {observed} before training/evaluating on this dataset." + ) + + @staticmethod + def _max_individuals_in_json(coco_json: dict) -> int: + """Returns the max number of annotations on any single image in a COCO dict.""" + img_to_annotations = map_id_to_annotations(coco_json.get("annotations") or []) + if not img_to_annotations: + return 0 + return max(len(ann_ids) for ann_ids in img_to_annotations.values()) + def get_dataset_parameters(self) -> PoseDatasetParameters: """Retrieves dataset parameters based on the instance's configuration. @@ -81,7 +124,9 @@ def get_dataset_parameters(self) -> PoseDatasetParameters: An instance of the PoseDatasetParameters with the parameters set. """ if self._dataset_parameters is None: - num_individuals, bodyparts = self.get_project_parameters(self.train_json) + meta = self.model_cfg.metadata + bodyparts = meta.bodyparts + individuals = meta.individuals crop_cfg = self.model_cfg.select("data.train.top_down_crop") or {} crop_w, crop_h = crop_cfg.get("width", 256), crop_cfg.get("height", 256) @@ -94,8 +139,8 @@ def get_dataset_parameters(self) -> PoseDatasetParameters: self._dataset_parameters = PoseDatasetParameters( bodyparts=bodyparts, - unique_bpts=[], - individuals=[f"individual{i}" for i in range(num_individuals)], + unique_bpts=meta.unique_bodyparts, + individuals=individuals, with_center_keypoints=self.model_cfg.get("with_center_keypoints", False), color_mode=self.model_cfg.get("color_mode", "RGB"), ctd_bbox_margin=ctd_bbox_margin, @@ -298,28 +343,43 @@ def load_data(self, mode: str = "train") -> dict: return data @staticmethod - def get_project_parameters(train_json: dict) -> tuple[int, list[str]]: - """ - Loads the parameters for the project from the train json file - TODO: Should this compute the number also using the test json? + def get_project_parameters( + train_json: dict, + test_json: dict | None = None, + ) -> tuple[int, list[str]]: + """Suggests parameters for a project, given its COCO-format JSON annotation(s). + + Use this to pick `bodyparts`/`max_individuals` when building a `PoseConfig` for + a new COCO project (e.g. before calling `make_pytorch_pose_config`). Once a + model config exists, it becomes authoritative for the dataset (see + `COCOLoader.get_dataset_parameters`) - this helper is only meant to bootstrap + it from the data. Args: train_json: the json dictionary containing the data for training + test_json: the json dictionary containing the data for testing/evaluation, + if any. Passing this ensures the suggested number of individuals also + covers the test set, so a model trained with it doesn't fail during + evaluation because a test image has more individuals than train ever did. Returns: - int: the maximum number of individuals in a single image + int: the maximum number of individuals in a single image, across train + (and test, if given) list[str]: the name of keypoints annotated in this project + + Raises: + ValueError: If the train JSON contains no images. """ - # TODO: Check that there's a single category + train_json = COCOLoader.validate_categories(train_json) bodyparts = train_json["categories"][0]["keypoints"] - img_to_annotations = map_id_to_annotations(train_json["annotations"]) - if len(img_to_annotations) == 0: + num_individuals = COCOLoader._max_individuals_in_json(train_json) + if num_individuals == 0: raise ValueError(f"No images found in the dataset: {train_json}!") - elif len(img_to_annotations) == 1: - num_individuals = len(list(img_to_annotations.values())[0]) - else: - num_individuals = max(*[len(a_ids) for a_ids in img_to_annotations.values()]) + + if test_json is not None: + test_json = COCOLoader.validate_categories(test_json) + num_individuals = max(num_individuals, COCOLoader._max_individuals_in_json(test_json)) return num_individuals, bodyparts diff --git a/dev-docs/docs/developer-guides/dataprep.md b/dev-docs/docs/developer-guides/dataprep.md index d72c787884..ffb20afb5c 100644 --- a/dev-docs/docs/developer-guides/dataprep.md +++ b/dev-docs/docs/developer-guides/dataprep.md @@ -79,12 +79,18 @@ import deeplabcut.pose_estimation_pytorch as dlc_torch project_root = Path("/path/to/COCOProject") -# Parse dataset information +# Parse dataset information (test.json is optional; include if you have one) train_dict = dlc_torch.COCOLoader.load_json( project_root, filename="train.json" ) -max_num_individuals, bodyparts = dlc_torch.COCOLoader.get_project_parameters(train_dict) +test_dict = dlc_torch.COCOLoader.load_json( + project_root, + filename="test.json" +) +max_num_individuals, bodyparts = dlc_torch.COCOLoader.get_project_parameters( + train_dict, test_dict +) # Create model configuration model_cfg = dlc_torch.config.make_pytorch_pose_config( diff --git a/tests/pose_estimation_pytorch/data/test_cocoloader.py b/tests/pose_estimation_pytorch/data/test_cocoloader.py new file mode 100644 index 0000000000..8900e7132c --- /dev/null +++ b/tests/pose_estimation_pytorch/data/test_cocoloader.py @@ -0,0 +1,251 @@ +# +# DeepLabCut Toolbox (deeplabcut.org) +# © A. & M.W. Mathis Labs +# https://github.com/DeepLabCut/DeepLabCut +# +# Please see AUTHORS for contributors. +# https://github.com/DeepLabCut/DeepLabCut/blob/main/AUTHORS +# +# Licensed under GNU Lesser General Public License v3.0 +# +"""Tests for COCOLoader dataset parameter handling.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +from PIL import Image + +from deeplabcut.core.config import ProjectConfig +from deeplabcut.pose_estimation_pytorch.config.pose import PoseConfig +from deeplabcut.pose_estimation_pytorch.data.cocoloader import COCOLoader + +BODYPARTS = ["nose", "tail"] + + +def _ann(image_id: int, ann_id: int, keypoints: list[float] | None = None) -> dict: + if keypoints is None: + # two bodyparts visible + keypoints = [10.0, 10.0, 2.0, 20.0, 20.0, 2.0] + return { + "id": ann_id, + "image_id": image_id, + "category_id": 1, + "keypoints": keypoints, + "bbox": [5.0, 5.0, 30.0, 30.0], + "num_keypoints": 2, + "iscrowd": 0, + "area": 900.0, + } + + +def _coco_dict( + image_name: str, + image_id: int, + n_individuals: int, + start_ann_id: int = 1, + bodyparts: list[str] | None = None, +) -> dict: + return { + "images": [ + { + "id": image_id, + "file_name": image_name, + "width": 64, + "height": 64, + } + ], + "annotations": [_ann(image_id, start_ann_id + i) for i in range(n_individuals)], + "categories": [ + { + "id": 1, + "name": "animal", + "keypoints": bodyparts or BODYPARTS, + "skeleton": [], + } + ], + } + + +def _write_project( + tmp_path: Path, + *, + train_n: int, + test_n: int | None, + max_individuals: int, + bodyparts: list[str] | None = None, + json_bodyparts: list[str] | None = None, +) -> tuple[Path, PoseConfig]: + """Writes a minimal COCO project (train + optional test json) and a matching PoseConfig. + + Args: + train_n: number of individuals annotated on the (single) train image. + test_n: number of individuals annotated on the (single) test image, or None to + skip writing a test.json altogether. + max_individuals: the number of individuals to configure in the PoseConfig. + bodyparts: the bodyparts to put in the PoseConfig (defaults to BODYPARTS). + json_bodyparts: the bodyparts to put in the COCO json category (defaults to + `bodyparts`, i.e. matching). Set to something else to simulate a mismatch. + """ + bodyparts = bodyparts or BODYPARTS + json_bodyparts = json_bodyparts if json_bodyparts is not None else bodyparts + + project_root = tmp_path / "coco_project" + ann_dir = project_root / "annotations" + img_dir = project_root / "images" + ann_dir.mkdir(parents=True) + img_dir.mkdir(parents=True) + + Image.new("RGB", (64, 64), color=(128, 128, 128)).save(img_dir / "train.png") + Image.new("RGB", (64, 64), color=(64, 64, 64)).save(img_dir / "test.png") + + train = _coco_dict("train.png", image_id=1, n_individuals=train_n, start_ann_id=1, bodyparts=json_bodyparts) + (ann_dir / "train.json").write_text(json.dumps(train)) + + if test_n is not None: + test = _coco_dict("test.png", image_id=2, n_individuals=test_n, start_ann_id=100, bodyparts=json_bodyparts) + (ann_dir / "test.json").write_text(json.dumps(test)) + + if max_individuals > 1: + project_config = ProjectConfig( + project_path=project_root, + bodyparts="MULTI!", + multianimalbodyparts=bodyparts, + individuals=[f"individual{i}" for i in range(max_individuals)], + multianimalproject=True, + ) + else: + project_config = ProjectConfig( + project_path=project_root, + bodyparts=bodyparts, + individuals=["individual0"], + multianimalproject=False, + ) + pose_config_path = project_root / "pytorch_config.yaml" + pose_config = PoseConfig.build( + project_config, + pose_config_path, + top_down=False, + net_type="resnet_50", + multi_animal=max_individuals > 1, + ) + return project_root, pose_config + + +def _make_loader(project_root: Path, pose_config: PoseConfig, has_test: bool) -> COCOLoader: + return COCOLoader( + project_root=project_root, + model_config=pose_config, + test_json_filename="test.json" if has_test else "", + ) + + +def test_max_individuals_in_json(): + coco = _coco_dict("a.png", image_id=1, n_individuals=3) + coco["annotations"].append(_ann(image_id=1, ann_id=99)) # 4 on same image + assert COCOLoader._max_individuals_in_json(coco) == 4 + assert COCOLoader._max_individuals_in_json({"annotations": []}) == 0 + + +def test_loader_accepts_capacity_covering_both_splits(tmp_path: Path): + # train max=3, test max=4, config capacity=4 -> should load without error, using + # the config's capacity (not the train json's). + project_root, pose_config = _write_project(tmp_path, train_n=3, test_n=4, max_individuals=4) + loader = _make_loader(project_root, pose_config, has_test=True) + + params = loader.get_dataset_parameters() + assert params.max_num_animals == 4 + assert list(params.individuals) == list(pose_config.metadata.individuals) + assert list(params.bodyparts) == BODYPARTS + + +def test_loader_raises_when_individuals_exceed_capacity(tmp_path: Path): + # train max=3, test max=4, config capacity=3 -> test.json needs more individuals + # than the model supports; this must fail loudly, at construction time. + project_root, pose_config = _write_project(tmp_path, train_n=3, test_n=4, max_individuals=3) + + with pytest.raises(ValueError, match=r"test\.json has an image with 4 individuals"): + _make_loader(project_root, pose_config, has_test=True) + + +def test_loader_raises_when_train_alone_exceeds_capacity(tmp_path: Path): + # No test.json: the check must still trigger for train.json. + project_root, pose_config = _write_project(tmp_path, train_n=4, test_n=None, max_individuals=3) + + with pytest.raises(ValueError, match=r"train\.json has an image with 4 individuals"): + _make_loader(project_root, pose_config, has_test=False) + + +def test_loader_raises_on_bodypart_mismatch(tmp_path: Path): + project_root, pose_config = _write_project( + tmp_path, + train_n=1, + test_n=1, + max_individuals=1, + bodyparts=["snout", "tailbase"], + json_bodyparts=BODYPARTS, # differs from the PoseConfig's bodyparts + ) + + with pytest.raises(ValueError, match="don't match model_cfg.metadata.bodyparts"): + _make_loader(project_root, pose_config, has_test=True) + + +def test_loader_ok_without_test_json(tmp_path: Path): + project_root, pose_config = _write_project(tmp_path, train_n=2, test_n=None, max_individuals=2) + loader = _make_loader(project_root, pose_config, has_test=False) + + params = loader.get_dataset_parameters() + assert params.max_num_animals == 2 + + +def test_get_project_parameters_train_only(): + train = _coco_dict("train.png", image_id=1, n_individuals=3) + num_individuals, bodyparts = COCOLoader.get_project_parameters(train) + assert num_individuals == 3 + assert list(bodyparts) == BODYPARTS + + +def test_get_project_parameters_considers_test_json(): + # see https://github.com/DeepLabCut/DeepLabCut/issues/3432 + train = _coco_dict("train.png", image_id=1, n_individuals=3) + test = _coco_dict("test.png", image_id=2, n_individuals=4) + + num_individuals, bodyparts = COCOLoader.get_project_parameters(train) + assert num_individuals == 3 + + num_individuals, bodyparts = COCOLoader.get_project_parameters(train, test) + assert num_individuals == 4 + assert list(bodyparts) == BODYPARTS + + +def test_get_project_parameters_raises_on_empty_train_json(): + empty = _coco_dict("a.png", image_id=1, n_individuals=0) + with pytest.raises(ValueError, match="No images found"): + COCOLoader.get_project_parameters(empty) + + +def test_get_project_parameters_warns_on_multiple_categories(): + train = _coco_dict("train.png", image_id=1, n_individuals=2) + train["categories"].append({"id": 2, "name": "other", "keypoints": ["eye"], "skeleton": []}) + + with pytest.warns(UserWarning, match="more than 1 category"): + num_individuals, bodyparts = COCOLoader.get_project_parameters(train) + + assert num_individuals == 2 + assert list(bodyparts) == BODYPARTS + + +def test_get_project_parameters_warns_on_multiple_categories_in_test_json(): + # The same category validation/normalization applied to train.json must also be + # applied to test.json, not skipped just because we don't read its bodyparts. + train = _coco_dict("train.png", image_id=1, n_individuals=2) + test = _coco_dict("test.png", image_id=2, n_individuals=2) + test["categories"].append({"id": 2, "name": "other", "keypoints": ["eye"], "skeleton": []}) + + with pytest.warns(UserWarning, match="more than 1 category"): + num_individuals, bodyparts = COCOLoader.get_project_parameters(train, test) + + assert num_individuals == 2 + assert list(bodyparts) == BODYPARTS