diff --git a/deeplabcut/pose_estimation_pytorch/apis/evaluation.py b/deeplabcut/pose_estimation_pytorch/apis/evaluation.py index f9b02e2c78..f4f0ec91c5 100755 --- a/deeplabcut/pose_estimation_pytorch/apis/evaluation.py +++ b/deeplabcut/pose_estimation_pytorch/apis/evaluation.py @@ -114,6 +114,7 @@ def evaluate( comparison_bodyparts: str | list[str] | None = None, per_keypoint_evaluation: bool = False, pcutoff: float | list[float] = 0.6, + force_multi_animal: bool = False, ) -> tuple[dict[str, float], dict[str, dict[str, np.ndarray]]]: """ Args: @@ -134,6 +135,9 @@ def evaluate( pcutoff: Confidence threshold for RMSE computation. If a list is provided, there should be one value for each bodypart and one value for each unique bodypart (if there are any). + force_multi_animal: If False - the scenario (single- or multi-animal) is inferred from the loader. + If True - the multi-animal is used during evaluation, even if the loader contains only a single animal. + Returns: A dict containing the evaluation results @@ -201,7 +205,7 @@ def evaluate( results = metrics.compute_metrics( gt_pose, pred_pose, - single_animal=parameters.max_num_animals == 1, + single_animal=False if force_multi_animal else parameters.max_num_animals == 1, pcutoff=pcutoff, unique_bodypart_poses=pred_unique, unique_bodypart_gt=gt_unique, diff --git a/deeplabcut/pose_estimation_pytorch/data/cocoloader.py b/deeplabcut/pose_estimation_pytorch/data/cocoloader.py index f086cfd4ee..1fc6ed6f05 100644 --- a/deeplabcut/pose_estimation_pytorch/data/cocoloader.py +++ b/deeplabcut/pose_estimation_pytorch/data/cocoloader.py @@ -29,6 +29,7 @@ class COCOLoader(Loader): """ Attributes: project_root: root directory path of the COCO project. + model_config_path: path to the pytorch_config.yaml file train_json_filename: the name of the json file containing the train annotations test_json_filename: the name of the json file containing the train annotations. None if there is no test set. @@ -79,7 +80,9 @@ def get_dataset_parameters(self) -> PoseDatasetParameters: bodyparts=bodyparts, unique_bpts=[], individuals=[f"individual{i}" for i in range(num_individuals)], - with_center_keypoints=self.model_cfg.get("with_center_keypoints", False), + with_center_keypoints=self.model_cfg.get( + "with_center_keypoints", False + ), color_mode=self.model_cfg.get("color_mode", "RGB"), top_down_crop_size=(crop_w, crop_h), top_down_crop_margin=crop_margin, @@ -262,15 +265,15 @@ def load_data(self, mode: str = "train") -> dict: annotations_per_image[image_id] = individual_idx + 1 filter_annotations = [] - for annotation in data['annotations']: - keypoints = annotation['keypoints'] - bbox = annotation['bbox'] + for annotation in data["annotations"]: + keypoints = annotation["keypoints"] + bbox = annotation["bbox"] if np.all(keypoints <= 0) or len(bbox) == 0: continue filter_annotations.append(annotation) - data["annotations"] = filter_annotations - + data["annotations"] = filter_annotations + # FIXME: why estimating bbox when there are already bbox? annotations_with_bbox = self._compute_bboxes( data["images"], diff --git a/deeplabcut/pose_estimation_pytorch/data/dlcloader.py b/deeplabcut/pose_estimation_pytorch/data/dlcloader.py index c8f43d007e..bd61c59f31 100644 --- a/deeplabcut/pose_estimation_pytorch/data/dlcloader.py +++ b/deeplabcut/pose_estimation_pytorch/data/dlcloader.py @@ -25,6 +25,7 @@ from deeplabcut.pose_estimation_pytorch.data.base import Loader from deeplabcut.pose_estimation_pytorch.data.dataset import PoseDatasetParameters from deeplabcut.pose_estimation_pytorch.data.snapshots import Snapshot +from deeplabcut.pose_estimation_pytorch.data.utils import bbox_from_keypoints from deeplabcut.pose_estimation_pytorch.data.utils import read_image_shape_fast @@ -481,7 +482,42 @@ def to_coco( } ) - return {"annotations": anns, "categories": categories, "images": images} + coco_dict = {"annotations": anns, "categories": categories, "images": images} + coco_dict = DLCLoader._add_bbox_annotations(coco_dict) + coco_dict = DLCLoader._remove_nans(coco_dict) + return coco_dict + + @staticmethod + def _add_bbox_annotations(coco_dict: dict) -> dict: + for annotation in coco_dict.get("annotations", []): + if "bbox" not in annotation: + image = [ + img + for img in coco_dict.get("images") + if img.get("id") == annotation.get("image_id") + ][0] + bbox = bbox_from_keypoints( + keypoints=np.array( + annotation["keypoints"] + ), # (..., num_keypoints, xy) + image_h=image.get("height"), + image_w=image.get("width"), + margin=20, + ) + annotation["bbox"] = list(bbox) + return coco_dict + + @staticmethod + def _remove_nans(coco_dict: dict) -> dict: + # Iterate through annotations and fix keypoints + for annotation in coco_dict.get("annotations", []): + if "keypoints" in annotation: + for keypoint in annotation["keypoints"]: + if any(isinstance(v, float) and np.isnan(v) for v in keypoint[:2]): + keypoint[0] = 0.0 # Replace x with 0 + keypoint[1] = 0.0 # Replace y with 0 + keypoint[2] = 0.0 # Ensure visibility is also 0 + return coco_dict @property def _dfs(self) -> dict[str, pd.DataFrame]: