Conversation
Lets training filter a dataset down to a subset of its declared classes
without editing annotation files: load_data_config(classes=[...]) builds
a {orig_id: orig_id} keep-set (identity, not compacted -- predictions and
checkpoint metadata stay directly comparable to the full dataset), threaded
through both label parsers and both dataset classes to drop boxes for
excluded ids. single_cls combined with classes= still collapses kept ids
to one merged class, since that is an intentional remap, not an exclusion.
BaseTrainer._setup_data() (and its two by-hand copies in DFINE/DEIM) build the filtered dataset via the new class_remap mechanism; nc/names resolve exactly as they would without classes=, since kept ids are not compacted. Also fixes a related, previously-existing bug: _rebuild_for_new_classes() always resets the wrapper to generic class_N placeholder names, and only single_cls got them restored afterward. Any other dataset-driven head resize (nc mismatch against the loaded checkpoint, classes= included) kept the placeholders. _resolve_num_classes_from_data_config() now stashes the dataset's real names for _sync_wrapped_model_num_classes to restore in every case. classes= is gated to G0/G1 detection families in _wrap_train_with_cfg, the same restriction as single_cls (both can trigger _rebuild_for_new_classes), and is inherited from the checkpoint on resume=True. DetectionValidator auto-inherits classes from the checkpoint's saved training config the same way single_cls already does, so val() needs nothing extra.
Mirrors --single-cls: classes=0,3,5 / --classes 0,3,5 (both grammars), gated to G0/G1 detection the same way, visible in --dry-run --json. Comma-separated string accepted directly by TrainConfig.classes, matching how device="0,1" is already written; no CLI-side parsing needed.
BaseTrainer._setup_data() already fell back to len(names) when a dataset yaml has no explicit nc key; _resolve_num_classes_from_data_config() and DFINE/DEIM's by-hand _setup_data() copies did not, instead keeping whatever num_classes the trainer started with (e.g. stale, from a previously loaded checkpoint). That produced a real class-count mismatch independent of classes=, surfaced as valid annotation ids being warned away as "out of range" against the wrong, smaller nc.
| else None | ||
| ), | ||
| single_cls=self._single_cls_enabled(), | ||
| classes=self.config.classes, |
There was a problem hiding this comment.
Subset Metrics Stay Unfiltered
Validation filters the dataloader annotations by classes, but _init_metrics separately builds COCO ground truth from the full JSON or YOLO labels, and _update_metrics submits predictions without filtering their classes. As a result, standalone and per-epoch validation report mAP over the full dataset rather than the requested subset. With native COCO data, excluded labels can also fall through the incomplete label-to-category mapping. Apply the selected original class IDs consistently to evaluator ground truth and predictions.
| classes=self.config.classes, | ||
| ) | ||
| class_remap = data_cfg.get("_class_remap") | ||
| if data_cfg.get("input_profile") is not None or getattr(self.wrapper_model, "input_profile", None) is not None: | ||
| from ..data.event_histogram import setup_histogram_data | ||
| return setup_histogram_data(self, data_cfg) |
There was a problem hiding this comment.
Histogram Training Ignores Classes
When a dataset uses an input_profile, this branch derives _class_remap and then immediately delegates to setup_histogram_data. That function creates YOLODataset without the remap. A supported G0/G1 run therefore accepts and saves classes, but histogram training still supervises every annotation class.
| # Auto-inherited from the checkpoint's saved training config the same | ||
| # way single_cls is (see DetectionValidator.__init__); rarely set by | ||
| # hand. Must match the classes= the model was trained with, since the | ||
| # head size is shared. | ||
| classes: Optional[List[int]] = field(default=None, kw_only=True) |
There was a problem hiding this comment.
Validation Classes Lack Checks
ValidationConfig adds classes without the normalization and validation already used by TrainConfig. Consequently, model.val(classes=[]) or a malformed checkpoint value can create an empty or negative remap and silently discard all dataloader labels, while the comma-separated form "0,3,5" fails during integer conversion. Apply the same string normalization and non-empty, non-negative, duplicate checks at validation entry points.
| classes: Optional[str] = typer.Option( | ||
| None, | ||
| help="Train a G0/G1 detector on only these original dataset class " | ||
| "ids, comma-separated (e.g. '0,3,5'); every other class is dropped " | ||
| "as if unlabeled. Ids are kept as-is, not compacted", | ||
| ), |
There was a problem hiding this comment.
Only the train CLI exposes the new dataset-subset option. libreyolo val ... classes=0,3,5 is still rejected because val_cmd declares no classes option and does not pass one to loaded_model.val. Checkpoint inheritance does not address users who explicitly want to evaluate a different subset.
| # --------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| def test_libreyolo9_train_end_to_end_keeps_original_nc_and_names(tmp_path): |
There was a problem hiding this comment.
This fast-unit test constructs a real LibreYOLO9 model, runs a complete training epoch, and writes a checkpoint. That adds integration-level runtime and resource sensitivity to routine unit execution. Move this coverage to an integration or smoke tier and retain a lightweight mocked API test here.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
|
Hello @gboeer, thanks for the follow up to the issue creating this PR. The overal approach is right but greptile flagged a few things which can be addressed in the next pass of code 🪖
|
This PR introduces a new feature as discussed in #828
What: Add
classes=to train()/val()/predict() (Python API + CLI) so a dataset's declared classes can be filtered to a subset without editing annotation files.Why: Users want to fine-tune on fewer classes than a dataset declares, without hand-editing labels or losing comparability with the original class ids.
classes=[...]filters which original class ids receive supervision; kept ids are NOT compacted to a contiguous range, so predictions/checkpoints stay directly comparable to the full dataset (generalizes the existingsingle_clsremap-at-parse-time mechanism)nc/namesare left untouched byclasses=; only which boxes reach the loss changesmodel.val()and the checkpoint resume path auto-inheritclassesfrom the saved training config, same pattern assingle_clslibreyolo train ... classes=0,3,5/--classes 0,3,5classes=into its ownValidationConfig, so periodic mAP was silently scored against the full, unfiltered datasetCode provenance
Original code written for this PR; bug fixes to LibreYOLO's own first-party code; no third-party code ported, adapted, or introduced; no GPL/AGPL/LGPL/non-commercial/unknown-license material involved.
The PR is not yet safe to merge because four previously reported behavioral gaps remain, including incomplete subset metrics and training paths that still ignore or cannot safely accept the option.
Summary
This PR adds class-subset filtering for G0/G1 detection training and validation while preserving original class IDs and the dataset’s full class metadata.
classesthrough configuration, dataset parsing, training, checkpoint resume, and periodic validation.Diagram
%%{init: {'theme': 'neutral'}}%% flowchart LR A[classes configuration] --> B[Build original-ID remap] B --> C[Training dataset] B --> D[Validation dataset] C --> E[Drop excluded annotations] D --> F[Filter evaluator ground truth] G[Model predictions] --> H[COCO evaluator] F --> HReviews (2) · Last reviewed commit: "fix: Subset Metrics Stay Unfiltered"