refactor(data): add keypoint detection data args for COCO and YOLO by liopeer · Pull Request #977 · lightly-ai/lightly-train

liopeer · 2026-09-21T09:22:40Z

First PR towards keypoint detection support (see NOTES_keypoint_detection_ecpose.md). It covers only the data args — the user-facing config classes and the keypoint set resolution. Reading labels off disk is the follow-up PR, stacked on this one.

What's in it

File	What
`_data/keypoint_helpers.py`	`KeypointSetArgs` + the pure resolvers that combine a dataset's keypoint set with a user-declared one
`_data/keypoint_detection_dataset.py`	`COCO`/`YOLOKeypointDetectionDataArgs`, mirroring `instance_segmentation_dataset.py`

Design decisions worth a look

No per-dataset presets. No COCO-17 or CrowdPose-14 preset, no baked sigma table, no baked flip pairs. The keypoint set comes from the dataset (COCO categories[], YOLO kpt_shape/flip_idx/kpt_names) or from the user's keypoints argument. Where a format carries nothing — COCO has no sigmas and no flip pairs, YOLO has no sigmas and no skeleton — the field stays None and whichever consumer needs it decides. Nothing is invented here.

Values given twice must agree. Neither the dataset nor the user takes precedence; a conflict is an error naming both values. Silent precedence in either direction would be impossible to debug.

One shared keypoint set. If several included COCO categories declare keypoints they must declare the same ones, and the same for YOLO's per-class kpt_names. A category that is excluded via ignore_classes may differ freely.

Keypoint set resolved at construction. A bad keypoint config fails when the args are built, not much later when labels are read. COCO categories[].skeleton is one-indexed and is converted to zero-indexed on read; flip_idx must be a permutation and its own inverse, since flipping twice has to be the identity.

Filtering fields ship here, enforced in the follow-up. min_keypoints, include_crowd and the skip_if_* flags are part of the data args schema and feed the mmap hashes, but the code that acts on them lives in the dataset args. Shipping the schema whole avoids changing it in the next PR.

min_keypoints defaults to 0. The reader reproduces the file faithfully and drops nothing. Filtering zero-keypoint instances is standard in pose training, but that's a training-policy call for the model PR. include_crowd=False is the one default that does drop data, so it's an explicit field.

Deliberately not wired up

No KeypointDetectionTrainTaskConfig, no keypoint_detection in TrainTaskConfig.task, no train_keypoint_detection(), no __init__.py export, no docs. That task field is a closed Literal; adding a member would advertise a task that TASK_TRAIN_MODEL_CLASSES cannot serve, so it would validate and then crash at model construction.

The format discriminator ships on both classes from day one, and TestKeypointDetectionDataConfig exercises the discriminated union plus the YAML path in the test file — that's where a field missing from the data args would be dropped silently by load_data_yaml_if_path. Those few lines move into train_task.py unchanged once a model lands. YOLOOrientedObjectDetectionDataArgs is the cautionary example of a data args class that shipped without format and so can't join a union.

Commit type

refactor rather than feat: per CONTRIBUTING.md, feat means a user-facing feature and bumps the minor version, and nothing here is reachable by users. feat belongs to the PR that ships train_keypoint_detection. Happy to change it if you'd rather this cut a release.

Verification

make format && make static-checks          # mypy clean over 650 files
uv run --frozen pytest tests/_data -q      # 369 passed

🤖 Generated with Claude Code

First PR towards keypoint detection support: the data args and dataset args that read a keypoint dataset off disk, plus tests. No model, no transforms, no train_task wiring. The keypoint set (count, names, sigmas, flip permutation, skeleton) is either read from the dataset or declared by the user, with no per-dataset presets baked in. Fields that neither source provides stay None so that the consumer decides, rather than this layer inventing a default. The on-disk format details were settled by inspecting real annotation files rather than from assumption; dev_tools/check_keypoint_detection_data_args.py records that inspection and is re-runnable. Notably, real num_dims == 2 YOLO pose labels use no unlabeled-keypoint sentinel at all, so the reader treats every keypoint of such a dataset as labeled instead of guessing one. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 60b5688f43

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

The script was a throwaway used once to settle the on-disk format questions before writing the readers. Committing it left ~270 lines that nothing runs, pointing at two external URLs that will rot, and pinned an Ultralytics download URL into the repo for no ongoing benefit. The findings it produced are what mattered, and those live where they are load-bearing: in the reader and data args docstrings. The verification run itself is recorded in the pull request. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Apply the new AGENTS.md documentation style to the keypoint detection files: bullet points over prose, shorter sentences, no filler. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

_iter_yolo_label_lines now yields (line number, line). The wrapper only existed so the three older callers could keep their signature; they unpack the number away instead. The number is in reach if those readers ever want file:line error messages too. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Keep this PR to the data args only. The dataset args, the dataset, the YOLO pose label reader and the keypoint dataset item move to a follow-up PR stacked on this one. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Drop the Attributes blocks. No other DataArgs class has one, and they documented an arbitrary subset of the fields. What is left is an attribute docstring on the five fields whose type does not already say what they hold. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

liopeer · 2026-09-21T14:14:45Z

/review

gabrielfruet

Small stuff, I will not block because of this but I highly recommend addressing

gabrielfruet · 2026-09-22T17:08:26Z

+        keypoint_helpers.validate_flip_idx(flip_idx, 3)
+
+
+def test_validate_kpt_names__invalid() -> None:


We only have the invalid variant here.

I think tests can also be used as documentation to understand the code. So I think it's worth it to add

gabrielfruet · 2026-09-22T17:12:04Z

+from lightly_train.types import PathLike
+
+
+class YOLOKeypointDetectionDataArgs(TaskDataArgs):


Why not separate files? we have coco_object_detection_dataset and yolo_object_detection_dataset

gabrielfruet · 2026-09-22T17:33:40Z

+    return kpt_shape
+
+
+def validate_flip_idx(flip_idx: list[int] | None, num_keypoints: int) -> None:


Although internal, I think it's worth it to add a better documentation on the functions here. At least was pretty hard to me to understand flip idx and sigmas.

chatgpt-codex-connector Bot reviewed Sep 21, 2026

View reviewed changes

Comment thread src/lightly_train/_data/keypoint_detection_dataset.py Outdated

Comment thread src/lightly_train/_data/keypoint_detection_dataset.py Outdated

liopeer and others added 5 commits September 21, 2026 10:02

docs(data): tighten keypoint docstrings and comments

aac65f5

Apply the new AGENTS.md documentation style to the keypoint detection files: bullet points over prose, shorter sentences, no filler. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

liopeer commented Sep 21, 2026

View reviewed changes

liopeer and others added 4 commits September 21, 2026 13:31

refactor(data): simplify keypoint data args

f1a710e

Merge branch 'main' into feat/keypoint-detection-data-args

7637986

docs(data): clarify keypoint data args

af506cc

test(data): standardize keypoint test names

232dfc2

liopeer mentioned this pull request Sep 21, 2026

refactor(data): add keypoint detection datasets #978

Open

2 of 4 tasks

liopeer added this pull request to stack #979 September 21, 2026 14:19

gabrielfruet approved these changes Sep 22, 2026

View reviewed changes

Navigation Menu

Repository navigation

refactor(data): add keypoint detection data args for COCO and YOLO - #977

refactor(data): add keypoint detection data args for COCO and YOLO#977
liopeer wants to merge 10 commits into
mainfrom
feat/keypoint-detection-data-args

liopeer commented Sep 21, 2026 •

edited

Loading

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

liopeer commented Sep 21, 2026

Uh oh!

gabrielfruet left a comment

Uh oh!

gabrielfruet Sep 22, 2026

Uh oh!

gabrielfruet Sep 22, 2026

Uh oh!

gabrielfruet Sep 22, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		keypoint_helpers.validate_flip_idx(flip_idx, 3)


		def test_validate_kpt_names__invalid() -> None:

		from lightly_train.types import PathLike


		class YOLOKeypointDetectionDataArgs(TaskDataArgs):

		return kpt_shape


		def validate_flip_idx(flip_idx: list[int] \| None, num_keypoints: int) -> None:

Conversation

liopeer commented Sep 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What's in it

Design decisions worth a look

Deliberately not wired up

Commit type

Verification

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

liopeer commented Sep 21, 2026

Uh oh!

gabrielfruet left a comment

Choose a reason for hiding this comment

Uh oh!

gabrielfruet Sep 22, 2026

Choose a reason for hiding this comment

Uh oh!

gabrielfruet Sep 22, 2026

Choose a reason for hiding this comment

Uh oh!

gabrielfruet Sep 22, 2026

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

liopeer commented Sep 21, 2026 •

edited

Loading