Visitar URL original
Improve `DLCLoader.to_coco()` and add force_multi_animal argument to … · DeepLabCut/DeepLabCut@12f3083 · GitHub
Skip to content

Commit 12f3083

Browse files
Improve DLCLoader.to_coco() and add force_multi_animal argument to dlc_torch.evaluate() (#2995)
* Add force_multi_animal argument to dlc_torch.evaluate() * Update docstring * Add bbox annotations and remove nans in DLCLoader.to_coco() * Black --------- Co-authored-by: Mackenzie Mathis <mathis@rowland.harvard.edu>
1 parent d9c856b commit 12f3083

3 files changed

Lines changed: 51 additions & 8 deletions

File tree

‎deeplabcut/pose_estimation_pytorch/apis/evaluation.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,7 @@ def evaluate(
114114
comparison_bodyparts: str | list[str] | None = None,
115115
per_keypoint_evaluation: bool = False,
116116
pcutoff: float | list[float] = 0.6,
117+
force_multi_animal: bool = False,
117118
) -> tuple[dict[str, float], dict[str, dict[str, np.ndarray]]]:
118119
"""
119120
Args:
@@ -134,6 +135,9 @@ def evaluate(
134135
pcutoff: Confidence threshold for RMSE computation. If a list is provided,
135136
there should be one value for each bodypart and one value for each unique
136137
bodypart (if there are any).
138+
force_multi_animal: If False - the scenario (single- or multi-animal) is inferred from the loader.
139+
If True - the multi-animal is used during evaluation, even if the loader contains only a single animal.
140+
137141
138142
Returns:
139143
A dict containing the evaluation results
@@ -201,7 +205,7 @@ def evaluate(
201205
results = metrics.compute_metrics(
202206
gt_pose,
203207
pred_pose,
204-
single_animal=parameters.max_num_animals == 1,
208+
single_animal=False if force_multi_animal else parameters.max_num_animals == 1,
205209
pcutoff=pcutoff,
206210
unique_bodypart_poses=pred_unique,
207211
unique_bodypart_gt=gt_unique,

‎deeplabcut/pose_estimation_pytorch/data/cocoloader.py‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ class COCOLoader(Loader):
2929
"""
3030
Attributes:
3131
project_root: root directory path of the COCO project.
32+
model_config_path: path to the pytorch_config.yaml file
3233
train_json_filename: the name of the json file containing the train annotations
3334
test_json_filename: the name of the json file containing the train annotations.
3435
None if there is no test set.
@@ -79,7 +80,9 @@ def get_dataset_parameters(self) -> PoseDatasetParameters:
7980
bodyparts=bodyparts,
8081
unique_bpts=[],
8182
individuals=[f"individual{i}" for i in range(num_individuals)],
82-
with_center_keypoints=self.model_cfg.get("with_center_keypoints", False),
83+
with_center_keypoints=self.model_cfg.get(
84+
"with_center_keypoints", False
85+
),
8386
color_mode=self.model_cfg.get("color_mode", "RGB"),
8487
top_down_crop_size=(crop_w, crop_h),
8588
top_down_crop_margin=crop_margin,
@@ -262,15 +265,15 @@ def load_data(self, mode: str = "train") -> dict:
262265
annotations_per_image[image_id] = individual_idx + 1
263266

264267
filter_annotations = []
265-
for annotation in data['annotations']:
266-
keypoints = annotation['keypoints']
267-
bbox = annotation['bbox']
268+
for annotation in data["annotations"]:
269+
keypoints = annotation["keypoints"]
270+
bbox = annotation["bbox"]
268271
if np.all(keypoints <= 0) or len(bbox) == 0:
269272
continue
270273
filter_annotations.append(annotation)
271274

272-
data["annotations"] = filter_annotations
273-
275+
data["annotations"] = filter_annotations
276+
274277
# FIXME: why estimating bbox when there are already bbox?
275278
annotations_with_bbox = self._compute_bboxes(
276279
data["images"],

‎deeplabcut/pose_estimation_pytorch/data/dlcloader.py‎

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
from deeplabcut.pose_estimation_pytorch.data.base import Loader
2626
from deeplabcut.pose_estimation_pytorch.data.dataset import PoseDatasetParameters
2727
from deeplabcut.pose_estimation_pytorch.data.snapshots import Snapshot
28+
from deeplabcut.pose_estimation_pytorch.data.utils import bbox_from_keypoints
2829
from deeplabcut.pose_estimation_pytorch.data.utils import read_image_shape_fast
2930

3031

@@ -481,7 +482,42 @@ def to_coco(
481482
}
482483
)
483484

484-
return {"annotations": anns, "categories": categories, "images": images}
485+
coco_dict = {"annotations": anns, "categories": categories, "images": images}
486+
coco_dict = DLCLoader._add_bbox_annotations(coco_dict)
487+
coco_dict = DLCLoader._remove_nans(coco_dict)
488+
return coco_dict
489+
490+
@staticmethod
491+
def _add_bbox_annotations(coco_dict: dict) -> dict:
492+
for annotation in coco_dict.get("annotations", []):
493+
if "bbox" not in annotation:
494+
image = [
495+
img
496+
for img in coco_dict.get("images")
497+
if img.get("id") == annotation.get("image_id")
498+
][0]
499+
bbox = bbox_from_keypoints(
500+
keypoints=np.array(
501+
annotation["keypoints"]
502+
), # (..., num_keypoints, xy)
503+
image_h=image.get("height"),
504+
image_w=image.get("width"),
505+
margin=20,
506+
)
507+
annotation["bbox"] = list(bbox)
508+
return coco_dict
509+
510+
@staticmethod
511+
def _remove_nans(coco_dict: dict) -> dict:
512+
# Iterate through annotations and fix keypoints
513+
for annotation in coco_dict.get("annotations", []):
514+
if "keypoints" in annotation:
515+
for keypoint in annotation["keypoints"]:
516+
if any(isinstance(v, float) and np.isnan(v) for v in keypoint[:2]):
517+
keypoint[0] = 0.0 # Replace x with 0
518+
keypoint[1] = 0.0 # Replace y with 0
519+
keypoint[2] = 0.0 # Ensure visibility is also 0
520+
return coco_dict
485521

486522
@property
487523
def _dfs(self) -> dict[str, pd.DataFrame]:

0 commit comments

Comments
 (0)