Cleanup, new lint rules, decreased complexity - #425
Draft
kozlov721 wants to merge 4 commits into
Draft
Conversation
kozlov721
requested review from
conorsim,
klemen1999 and
tersekmatija
and removed request for
a team
July 26, 2026 20:55
github-actions
Bot
added
documentation
Contributor
📝 WalkthroughWalkthroughThe pull request decomposes model configuration, training, export, inference, data loading, visualization, annotation, and utility logic into focused helpers. It also updates lint hooks, documentation examples, and integration-test setup while preserving existing workflows. ChangesCross-cutting refactor
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 16
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
luxonis_train/config/config.py (1)
146-156: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
last_body_indexcan be-1when the first node is a head.At Line 147, if node 0 has
"Head"in its name,last_body_indexbecomes-1, and every subsequent node without explicitinputsresolvesnames[-1], which is the node's own just-appended name → a self-referencing edge (later rejected by the acyclic check with a confusing message). Guard againsti == 0.🐛 Proposed fix
- if "Head" in name and last_body_index is None: + if "Head" in name and last_body_index is None and i > 0: last_body_index = i - 1🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/config/config.py` around lines 146 - 156, Guard the last_body_index assignment in the node-processing logic so a head at index 0 does not set it to -1. Update the condition around the "Head" check to require i > 0, preserving the existing previous-node resolution for later head nodes and preventing subsequent nodes from using names[-1] as their predecessor.
🧹 Nitpick comments (13)
pyproject.toml (1)
113-113: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winScope the
PLR0917ignore narrowly.Ignoring this rule globally can hide future functions with excessive positional arguments—the kind of complexity regression this PR is intended to expose. If existing APIs require exceptions, use per-file or line-level ignores with a brief rationale.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pyproject.toml` at line 113, Replace the global PLR0917 ignore in the Ruff configuration with narrowly scoped per-file or line-level exceptions only where existing APIs require them, and add a brief rationale for each exception. Ensure new functions remain subject to PLR0917.luxonis_train/lightning/luxonis_lightning.py (3)
1325-1334: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReuse
_output_orderinforward.
forward(Lines 256-263) builds the identical sorted(node_name, output_name, i)list inline; now that the helper exists, calling it there removes the duplicate ordering logic and keeps export/forward ordering guaranteed in sync.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/lightning/luxonis_lightning.py` around lines 1325 - 1334, Update forward to call the existing _output_order helper instead of constructing the sorted (node_name, output_name, index) list inline. Remove the duplicated ordering logic while preserving forward’s current output ordering and downstream behavior.
1212-1215: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDead loop:
for _ in sorted(val_eval_epochs)adds the same key repeatedly.The val metric key doesn't depend on the epoch, so the loop only reduces to "add if
val_eval_epochsis non-empty". Making that explicit removes the pointless sort/iteration and preserves the empty-set semantics.♻️ Proposed simplification
else: - for _ in sorted(val_eval_epochs): - metric_keys.add(f"val/metric/{formatted_node_name}/{sub_name}") + if val_eval_epochs: + metric_keys.add(f"val/metric/{formatted_node_name}/{sub_name}") metric_keys.add(f"test/metric/{formatted_node_name}/{sub_name}")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/lightning/luxonis_lightning.py` around lines 1212 - 1215, In the metric-key construction branch, replace the epoch loop around val_eval_epochs with a single conditional that adds the val metric key only when val_eval_epochs is non-empty. Keep the test metric key added unconditionally in the existing else branch, and remove the unnecessary sorting and iteration.
363-423: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueLoose
Anyannotations introduced by the helper extraction. The extracted helpers dropped the concrete types that are already known at the call sites (NodeWrapperfromNodes.traverse()/Nodes.items(), andlist[str]order lists per_get_node_order_mapping), so the new boundaries are untyped.
luxonis_train/lightning/luxonis_lightning.py#L363-L423: annotate thenodeparameter of_collect_node_results,_collect_losses,_update_metrics, and_collect_visualizationsasNodeWrapper(already imported at Line 40).luxonis_train/lightning/luxonis_lightning.py#L689-L691: type_prepare_checkpoint's return astuple[dict[str, Any], Version, list[str] | None, list[str]]and propagatelist[str] | None/list[str]to theold_order/new_orderparameters of_load_node_checkpointand_load_with_order_mapping.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/lightning/luxonis_lightning.py` around lines 363 - 423, Replace the loose node annotations in _collect_node_results, _collect_losses, _update_metrics, and _collect_visualizations with NodeWrapper. In luxonis_train/lightning/luxonis_lightning.py:363-423, update only these helper parameters; in luxonis_train/lightning/luxonis_lightning.py:689-691, annotate _prepare_checkpoint as returning tuple[dict[str, Any], Version, list[str] | None, list[str]] and use list[str] | None for old_order and list[str] for new_order in _load_node_checkpoint and _load_with_order_mapping.luxonis_train/config/config.py (1)
284-311: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueRenamed aliases are not registered in
names.Line 311 adds the original
namerather than the (possibly rewritten)module.alias, so a later module whose name equals a generated alias won't be detected as a duplicate. Trackingmodule.alias or module.namekeeps the set consistent with what is actually emitted.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/config/config.py` around lines 284 - 311, Update _make_node_module_names_unique so the names set records the module’s final emitted identifier after any alias generation or renaming, using module.alias or module.name rather than the original name variable. Preserve the existing duplicate detection and warning behavior.luxonis_train/core/core.py (1)
397-411: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winWeight normalization/download duplicated with
_normalize_weights_and_load_ckpt.
__init__already normalized the dict and downloaded the checkpoint; this re-does both (a second download viaLuxonisFileSystem.download). Passing the already-loadedckptthrough would avoid the redundant network/disk work.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/core/core.py` around lines 397 - 411, Update _load_initial_weights to reuse the already normalized and downloaded checkpoint produced by _normalize_weights_and_load_ckpt instead of rewrapping dicts or calling LuxonisFileSystem.download again. Preserve the existing config-weight warning and pass the reused checkpoint directly to lightning_module.load_checkpoint.luxonis_train/core/utils/annotate_utils.py (2)
102-109: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueType
headmore precisely.
head: objectloses all type information at the call site. The argument is alwayslt_module.nodes[head_name].module, so annotate it as the node/module base type (e.g.BaseNodeornn.Module) while keeping theisinstance(head, lxt.BaseHead)narrowing.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/core/utils/annotate_utils.py` around lines 102 - 109, Update the head parameter annotation in _annotated_records from object to the appropriate node/module base type used by lt_module.nodes[head_name].module, such as BaseNode or nn.Module, while preserving the existing isinstance(head, lxt.BaseHead) narrowing and behavior.
113-118: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winNarrow the
suppressscope to record construction only.
yieldinsidewith suppress(Exception)also swallows exceptions thrown back into the generator by the consumer (e.g. failures insideLuxonisDataset.add), which can silently drop errors unrelated to invalid predictions. Build the record inside the guard and yield outside it.♻️ Proposed refactor
- # Skip predictions that are invalid, e.g. outside the clipping range. - with suppress(Exception): - yield DatasetRecord(**record) + # Skip predictions that are invalid, e.g. outside the clipping range. + try: + dataset_record = DatasetRecord(**record) + except Exception: # noqa: BLE001 + continue + yield dataset_record🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/core/utils/annotate_utils.py` around lines 113 - 118, Update the record conversion path after the DatasetRecord check so the suppress(Exception) context covers only DatasetRecord(**record) construction; store the constructed record, exit the suppression scope, then yield it outside the guard. Preserve skipping invalid predictions while allowing consumer-thrown exceptions from the generator yield to propagate.luxonis_train/utils/annotation.py (2)
127-180: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low value
instance_segmentationandsegmentationshare thenorm_maskskey.If a head requires both labels, the second branch overwrites the first and both annotators emit the same masks. Use distinct keys (
norm_instance_masks/norm_semantic_masks) to make the helper collision-free.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/utils/annotation.py` around lines 127 - 180, Update _prepare_transformed so instance_segmentation stores transformed masks under norm_instance_masks and segmentation stores them under norm_semantic_masks, preventing either branch from overwriting the other when both labels are required. Keep the existing transformation logic unchanged.
301-314: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNon-deterministic annotation ordering.
required_labelsis aset, so_emit_annotationsemits per-task records in arbitrary order across runs. Iterating_ANNOTATORSand filtering by membership gives a stable order at no cost.♻️ Proposed refactor
- for task in required_labels: - annotator = _ANNOTATORS.get(task) - if annotator is not None: - yield from annotator( - head, img_path, preds_for_image, transformed, i - ) + for task, annotator in _ANNOTATORS.items(): + if task in required_labels: + yield from annotator( + head, img_path, preds_for_image, transformed, i + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/utils/annotation.py` around lines 301 - 314, Update _emit_annotations to iterate through _ANNOTATORS in its defined order and emit annotations only for tasks present in required_labels, instead of iterating required_labels directly. Preserve the existing annotator lookup and yield behavior for matching tasks.luxonis_train/utils/general.py (1)
365-388: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueTighten the helper return types and the
batch_sizefalsy check.Both helpers return
Iterator[Any], which drops theIterator[tuple[Tensor, ...]] | Iterator[Tensor]contract the public overloads advertise. Also,batch_size or int(bboxes[:, 0].max()) + 1silently falls back forbatch_size == 0; an explicitis Nonecheck matches_empty_batch_instances.♻️ Proposed refactor
-def _batched_instances( - bboxes: Tensor, args: tuple[Tensor, ...], batch_size: int | None -) -> Iterator[Any]: - n_batches = batch_size or int(bboxes[:, 0].max()) + 1 +def _batched_instances( + bboxes: Tensor, args: tuple[Tensor, ...], batch_size: int | None +) -> Iterator[tuple[Tensor, ...] | Tensor]: + n_batches = ( + batch_size + if batch_size is not None + else int(bboxes[:, 0].max()) + 1 + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/utils/general.py` around lines 365 - 388, Update _empty_batch_instances and _batched_instances to return Iterator[tuple[Tensor, ...]] | Iterator[Tensor], matching the public overload contract instead of Iterator[Any]. In _batched_instances, replace the falsy batch_size fallback with an explicit None check so batch_size == 0 produces zero batches while only None derives the count from bboxes.luxonis_train/utils/boundingbox.py (2)
343-384: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueRecompute of
x[candidate_mask_i]in the additional-fields path.Line 375 re-slices
x[candidate_mask_i]even though the same slice was already produced at Line 359. Hold the original slice in a local before_apply_objectnessmutates/replacescurr_outand reuse it.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/utils/boundingbox.py` around lines 343 - 384, In _nms_single_image, preserve the initial x[candidate_mask_i] slice in a local variable before _apply_objectness updates curr_out, then reuse that variable when selecting additional fields instead of re-slicing x[candidate_mask_i].
299-299: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win
keep_maskis allocated on CPU whilebboxes/curr_outmay be on CUDA.
torch.zeros(bboxes.size(0)).bool()has nodevice=, sobboxes[keep_mask]mixes devices. Adddevice=bboxes.devicewhile the code is being extracted.🛡️ Proposed fix
- keep_mask = torch.zeros(bboxes.size(0)).bool() + keep_mask = torch.zeros(bboxes.size(0), dtype=torch.bool, device=bboxes.device)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/utils/boundingbox.py` at line 299, Update the keep_mask allocation in the bounding-box filtering logic to create the boolean tensor on bboxes.device, ensuring it matches bboxes and curr_out when they are on CUDA. Preserve the existing size and boolean dtype.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@luxonis_train/attached_modules/base_attached_module.py`:
- Around line 223-260: Update _validate_parameter_types to handle signature
parameters absent from kwargs, preserving callee defaults instead of indexing
unconditionally; skip validation for missing keys while continuing to validate
every present value, so _add_parameter’s non-empty-default path cannot raise
KeyError.
In `@luxonis_train/attached_modules/visualizers/utils.py`:
- Around line 432-436: Update the ValueError message in the resize_along
validation to interpolate the actual resize_along value by making the string an
f-string. Preserve the existing validation condition and valid-options text.
In `@luxonis_train/callbacks/fail_on_no_train_batches.py`:
- Around line 55-62: Update _loader_details to retrieve
trainer.fit_loop._combined_loader via getattr with a None fallback, then
preserve the existing flattened-list handling while safely treating a missing
combined loader as unavailable. Ensure this diagnostic path cannot raise
AttributeError and still allows the caller to produce the actionable
zero-training-batches RuntimeError.
In `@luxonis_train/config/config.py`:
- Around line 1027-1032: Update the accumulate_grad_batches calculation in the
trainer configuration to clamp the computed value to at least 1, including when
trainer.batch_size exceeds 64. Preserve the existing assignment and logging
behavior while ensuring downstream loss weights and scheduling never receive
zero.
In `@luxonis_train/core/core.py`:
- Line 182: Update the torch.load call in the checkpoint-loading flow to pass
weights_only=True, preserving subsequent access to the plain-data config and
dataset_metadata fields; if that payload cannot support weights_only, correct
the suppression to # nosemgrep and include a justification.
- Around line 634-641: Update the YAML export flow around
_upload_export_artifact so the file is fully closed before uploading: retain the
generated path while inside the with open block, then invoke the upload
afterward. Pass the intended artifact name rather than f.name as the upload
name, avoiding the absolute file path.
- Around line 167-183: Update _normalize_weights_and_load_ckpt to report the
requested weights value in the safe_download failure message instead of
downloaded, and add an explicit fallback for unsupported weights types that
raises a clear validation error rather than returning None.
In `@luxonis_train/lightning/luxonis_lightning.py`:
- Around line 942-955: Update _prepare_balanced_labels to prune the segmentation
background only when the corresponding classification key exists in labels_copy;
skip segmentation-only entries without raising a missing-key error, while
preserving the existing slicing behavior for present classification labels.
In `@luxonis_train/nodes/base_node.py`:
- Around line 725-741: Add a catch-all branch to the attach_index pattern match
in get_attached so unsupported values raise ValueError instead of falling
through and returning None. Preserve the existing handling for "all", integer
indices, slices, and None, and include the invalid attach_index value in the
error context.
- Around line 619-622: Update the index calculation near input_name and the
name-in-xyz check so the regex-derived idx assignment runs only when the xyz
branch did not assign an index. Preserve indices 0, 1, and 2 for x, y, and z
respectively, while retaining the existing regex fallback for other parameter
names.
In `@luxonis_train/upgrade.py`:
- Around line 130-135: Update the config format check in the
configuration-loading function to compare Path.suffix against the dotted JSON
extension, ensuring .json files use json.loads while other files continue
through yaml.safe_load. Apply the same correction to the config handling in the
__main__.py config command.
In `@luxonis_train/utils/annotation.py`:
- Around line 80-82: Update default_annotate so an image with all-empty
predictions yields its bare annotation and then continues processing the
remaining image paths in the batch. Replace the early generator termination
after _is_all_empty(preds_for_image, required_labels) with iteration
continuation, preserving the existing annotation output.
- Around line 259-288: Update _annotate_classification and _annotate_text to use
the already-selected per-image values in transformed without indexing them by i
again. Compute the class from transformed["pred_classes"].argmax(), and read the
decoded text from transformed["pred_text"] as the list[tuple[str, float]]
result, preserving the existing annotation structure.
In `@luxonis_train/utils/boundingbox.py`:
- Around line 581-591: Update the fallback mask in _bbox_weight to create
torch.ones_like(target_scores.sum(-1)) with dtype=torch.bool, ensuring
torch.masked_select receives a boolean mask when mask_positive is absent.
In `@tests/integration/test_export_unique_identifiers.py`:
- Around line 20-36: Update the test input generation around dtype_map and the
loop over onnx_model.graph.input so supported ONNX element types receive
compatible generated values; for any unsupported type, raise a clear error
instead of using the current np.float32 fallback. Preserve the existing
dynamic-dimension handling and initializer exclusion.
In `@tests/integration/test_predefined_models.py`:
- Around line 48-49: Update the output assertion in _assert_infer_output so it
runs for the "loader" subtest, matching the value passed by the loop, while
preserving the existing dataset-case behavior as appropriate.
---
Outside diff comments:
In `@luxonis_train/config/config.py`:
- Around line 146-156: Guard the last_body_index assignment in the
node-processing logic so a head at index 0 does not set it to -1. Update the
condition around the "Head" check to require i > 0, preserving the existing
previous-node resolution for later head nodes and preventing subsequent nodes
from using names[-1] as their predecessor.
---
Nitpick comments:
In `@luxonis_train/config/config.py`:
- Around line 284-311: Update _make_node_module_names_unique so the names set
records the module’s final emitted identifier after any alias generation or
renaming, using module.alias or module.name rather than the original name
variable. Preserve the existing duplicate detection and warning behavior.
In `@luxonis_train/core/core.py`:
- Around line 397-411: Update _load_initial_weights to reuse the already
normalized and downloaded checkpoint produced by
_normalize_weights_and_load_ckpt instead of rewrapping dicts or calling
LuxonisFileSystem.download again. Preserve the existing config-weight warning
and pass the reused checkpoint directly to lightning_module.load_checkpoint.
In `@luxonis_train/core/utils/annotate_utils.py`:
- Around line 102-109: Update the head parameter annotation in
_annotated_records from object to the appropriate node/module base type used by
lt_module.nodes[head_name].module, such as BaseNode or nn.Module, while
preserving the existing isinstance(head, lxt.BaseHead) narrowing and behavior.
- Around line 113-118: Update the record conversion path after the DatasetRecord
check so the suppress(Exception) context covers only DatasetRecord(**record)
construction; store the constructed record, exit the suppression scope, then
yield it outside the guard. Preserve skipping invalid predictions while allowing
consumer-thrown exceptions from the generator yield to propagate.
In `@luxonis_train/lightning/luxonis_lightning.py`:
- Around line 1325-1334: Update forward to call the existing _output_order
helper instead of constructing the sorted (node_name, output_name, index) list
inline. Remove the duplicated ordering logic while preserving forward’s current
output ordering and downstream behavior.
- Around line 1212-1215: In the metric-key construction branch, replace the
epoch loop around val_eval_epochs with a single conditional that adds the val
metric key only when val_eval_epochs is non-empty. Keep the test metric key
added unconditionally in the existing else branch, and remove the unnecessary
sorting and iteration.
- Around line 363-423: Replace the loose node annotations in
_collect_node_results, _collect_losses, _update_metrics, and
_collect_visualizations with NodeWrapper. In
luxonis_train/lightning/luxonis_lightning.py:363-423, update only these helper
parameters; in luxonis_train/lightning/luxonis_lightning.py:689-691, annotate
_prepare_checkpoint as returning tuple[dict[str, Any], Version, list[str] |
None, list[str]] and use list[str] | None for old_order and list[str] for
new_order in _load_node_checkpoint and _load_with_order_mapping.
In `@luxonis_train/utils/annotation.py`:
- Around line 127-180: Update _prepare_transformed so instance_segmentation
stores transformed masks under norm_instance_masks and segmentation stores them
under norm_semantic_masks, preventing either branch from overwriting the other
when both labels are required. Keep the existing transformation logic unchanged.
- Around line 301-314: Update _emit_annotations to iterate through _ANNOTATORS
in its defined order and emit annotations only for tasks present in
required_labels, instead of iterating required_labels directly. Preserve the
existing annotator lookup and yield behavior for matching tasks.
In `@luxonis_train/utils/boundingbox.py`:
- Around line 343-384: In _nms_single_image, preserve the initial
x[candidate_mask_i] slice in a local variable before _apply_objectness updates
curr_out, then reuse that variable when selecting additional fields instead of
re-slicing x[candidate_mask_i].
- Line 299: Update the keep_mask allocation in the bounding-box filtering logic
to create the boolean tensor on bboxes.device, ensuring it matches bboxes and
curr_out when they are on CUDA. Preserve the existing size and boolean dtype.
In `@luxonis_train/utils/general.py`:
- Around line 365-388: Update _empty_batch_instances and _batched_instances to
return Iterator[tuple[Tensor, ...]] | Iterator[Tensor], matching the public
overload contract instead of Iterator[Any]. In _batched_instances, replace the
falsy batch_size fallback with an explicit None check so batch_size == 0
produces zero batches while only None derives the count from bboxes.
In `@pyproject.toml`:
- Line 113: Replace the global PLR0917 ignore in the Ruff configuration with
narrowly scoped per-file or line-level exceptions only where existing APIs
require them, and add a brief rationale for each exception. Ensure new functions
remain subject to PLR0917.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 98373d48-1b44-4875-9d8a-ecbdd7a21733
📒 Files selected for processing (37)
.pre-commit-config.yamlREADME.mdexamples/README.mdluxonis_train/__main__.pyluxonis_train/attached_modules/base_attached_module.pyluxonis_train/attached_modules/metrics/base_metric.pyluxonis_train/attached_modules/visualizers/fomo_visualizer.pyluxonis_train/attached_modules/visualizers/utils.pyluxonis_train/callbacks/ema.pyluxonis_train/callbacks/fail_on_no_train_batches.pyluxonis_train/callbacks/luxonis_progress_bar.pyluxonis_train/config/config.pyluxonis_train/config/predefined_models/base_predefined_model.pyluxonis_train/core/core.pyluxonis_train/core/utils/annotate_utils.pyluxonis_train/core/utils/export_utils.pyluxonis_train/core/utils/infer_utils.pyluxonis_train/core/utils/tune_utils.pyluxonis_train/lightning/luxonis_lightning.pyluxonis_train/lightning/utils.pyluxonis_train/loaders/base_loader.pyluxonis_train/loaders/dummy_loader.pyluxonis_train/loaders/luxonis_loader_torch.pyluxonis_train/nodes/backbones/rexnetv1.pyluxonis_train/nodes/base_node.pyluxonis_train/nodes/blocks/__init__.pyluxonis_train/nodes/blocks/blocks.pyluxonis_train/nodes/blocks/unet.pyluxonis_train/nodes/necks/reppan_neck/reppan_neck.pyluxonis_train/upgrade.pyluxonis_train/utils/annotation.pyluxonis_train/utils/boundingbox.pyluxonis_train/utils/general.pypyproject.tomltests/integration/test_bump_opset_version.pytests/integration/test_export_unique_identifiers.pytests/integration/test_predefined_models.py
💤 Files with no reviewable changes (1)
- luxonis_train/nodes/blocks/init.py
Comment on lines
+259
to
+288
| def _annotate_classification( | ||
| head: "lxt.nodes.BaseHead", | ||
| img_path: Path, | ||
| preds_for_image: dict[str, Tensor], | ||
| transformed: dict[str, Any], | ||
| i: int, | ||
| ) -> DatasetIterator: | ||
| yield { | ||
| "file": str(img_path), | ||
| "task_name": head.task_name, | ||
| "annotation": { | ||
| "class": head.classes.inverse[ | ||
| int(transformed["pred_classes"][i].argmax()) | ||
| ], | ||
| }, | ||
| } | ||
|
|
||
|
|
||
| def _annotate_text( | ||
| head: "lxt.nodes.BaseHead", | ||
| img_path: Path, | ||
| preds_for_image: dict[str, Tensor], | ||
| transformed: dict[str, Any], | ||
| i: int, | ||
| ) -> DatasetIterator: | ||
| yield { | ||
| "file": str(img_path), | ||
| "task_name": head.task_name, | ||
| "annotation": {"metadata": {"text": transformed["pred_text"][i][0]}}, | ||
| } |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n --type=py -C5 'pred_classes|pred_text' luxonis_train/
rg -n --type=py -C3 'def decoder|self\.decoder' luxonis_train/Repository: luxonis/luxonis-train
Length of output: 22477
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== annotation.py relevant sections =="
sed -n '130,290p' luxonis_train/utils/annotation.py
echo
echo "== ocr_ctc_head.py relevant sections =="
sed -n '1,180p' luxonis_train/nodes/heads/ocr_ctc_head.py
echo
echo "== search _build_preds_for_image definition and call sites =="
rg -n --type=py -C8 '_build_preds_for_image|_ANNOTATORS|_annotate_classification|_annotate_text' luxonis_train/utils/annotation.pyRepository: luxonis/luxonis-train
Length of output: 12356
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== OCRDecoder definition =="
rg -n --type=py -C20 'class OCRDecoder|def _.*decode|ctc|text' luxonis_train -g '*.py' | sed -n '1,220p'
echo
echo "== behavioral probe for shapes/index semantics =="
python3 - <<'PY'
import numpy as np
i = 2
pred_classes = np.arange(5 * i + 3).reshape(-1) # shape [13]; i=2 selects one logit
print("classification shape", pred_classes.shape)
print("current classification annotation idx", int(pred_classes[i].argmax()), "from logit", pred_classes[i])
print("intended classification annotation idx", int(pred_classes.argmax()), "from logits")
decoded_texts = ["abc", "def", "ghi", "jkl"] # decoded output already has batch length 4
print("text shape", np.array(decoded_texts).shape)
print("current text annotation", decoded_texts[i][0] if i < len(decoded_texts) else "<out-of-bounds>")
print("intended text annotation", decoded_texts[0].split()[0] if False else None)
PYRepository: luxonis/luxonis-train
Length of output: 19527
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate OCRDecoder =="
rg -n 'OCRDecoder|class .*Decoder|decode\(self' luxonis_train -g '*.py'
echo
echo "== inspect candidate files =="
files=$(rg -l 'OCRDecoder|decode\(self' luxonis_train -g '*.py')
for f in $files; do
echo "--- $f ---"
wc -l "$f"
rg -n -C25 'class .*Decoder|def decode|decode\(self|class OCRDecoder|CTC|text' "$f"
doneRepository: luxonis/luxonis-train
Length of output: 4881
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== annotation.py full outline/top =="
sed -n '1,130p' luxonis_train/utils/annotation.py
echo
echo "== luxonis_train/utils/ocr.py =="
cat -n luxonis_train/utils/ocr.pyRepository: luxonis/luxonis-train
Length of output: 9525
Use the per-image quantities in the annotation helpers.
_build_preds_for_image() already selects head_output[task][i], so transformed["pred_classes"] is the classification logits for one image and transformed["pred_text"] is a list[tuple[str, float]], not a per-row tensor. Indexing again with i produces the wrong classification class and the wrong decoded text for every image except the first, or can crash when the decoded OCR result contains fewer decoded steps than the batch index.
🐛 Suggested direction
- "class": head.classes.inverse[
- int(transformed["pred_classes"][i].argmax())
- ],
+ "class": head.classes.inverse[
+ int(transformed["pred_classes"].argmax())
+ ],- "annotation": {"metadata": {"text": transformed["pred_text"][i][0]}},
+ "annotation": {"metadata": {"text": transformed["pred_text"][0]}},📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def _annotate_classification( | |
| head: "lxt.nodes.BaseHead", | |
| img_path: Path, | |
| preds_for_image: dict[str, Tensor], | |
| transformed: dict[str, Any], | |
| i: int, | |
| ) -> DatasetIterator: | |
| yield { | |
| "file": str(img_path), | |
| "task_name": head.task_name, | |
| "annotation": { | |
| "class": head.classes.inverse[ | |
| int(transformed["pred_classes"][i].argmax()) | |
| ], | |
| }, | |
| } | |
| def _annotate_text( | |
| head: "lxt.nodes.BaseHead", | |
| img_path: Path, | |
| preds_for_image: dict[str, Tensor], | |
| transformed: dict[str, Any], | |
| i: int, | |
| ) -> DatasetIterator: | |
| yield { | |
| "file": str(img_path), | |
| "task_name": head.task_name, | |
| "annotation": {"metadata": {"text": transformed["pred_text"][i][0]}}, | |
| } | |
| def _annotate_classification( | |
| head: "lxt.nodes.BaseHead", | |
| img_path: Path, | |
| preds_for_image: dict[str, Tensor], | |
| transformed: dict[str, Any], | |
| i: int, | |
| ) -> DatasetIterator: | |
| yield { | |
| "file": str(img_path), | |
| "task_name": head.task_name, | |
| "annotation": { | |
| "class": head.classes.inverse[ | |
| int(transformed["pred_classes"].argmax()) | |
| ], | |
| }, | |
| } | |
| def _annotate_text( | |
| head: "lxt.nodes.BaseHead", | |
| img_path: Path, | |
| preds_for_image: dict[str, Tensor], | |
| transformed: dict[str, Any], | |
| i: int, | |
| ) -> DatasetIterator: | |
| yield { | |
| "file": str(img_path), | |
| "task_name": head.task_name, | |
| "annotation": {"metadata": {"text": transformed["pred_text"][0]}}, | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@luxonis_train/utils/annotation.py` around lines 259 - 288, Update
_annotate_classification and _annotate_text to use the already-selected
per-image values in transformed without indexing them by i again. Compute the
class from transformed["pred_classes"].argmax(), and read the decoded text from
transformed["pred_text"] as the list[tuple[str, float]] result, preserving the
existing annotation structure.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #425 +/- ##
===========================================
- Coverage 93.70% 61.93% -31.78%
===========================================
Files 269 268 -1
Lines 13178 13490 +312
===========================================
- Hits 12349 8355 -3994
- Misses 829 5135 +4306
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
luxonis_train/core/core.py (2)
707-713: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winSerialize
input_modelas a string before writing YAML.
_build_modelconverter_configstores the exportedPathdirectly in the config and passes it toyaml.safe_dump; PyYAML cannot representPath, so this raisesRepresenterErrorand prevents the.yamlexport artifact from being written.🐛 Proposed fix
- "input_model": onnx_save_path, + "input_model": str(onnx_save_path),🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/core/core.py` around lines 707 - 713, Update _build_modelconverter_config so the input_model value is converted to a string before the returned configuration is passed to YAML serialization, while preserving the existing path value and all other configuration fields.
1708-1721: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winLoad the requested checkpoint before reparameterizing the quantization model.
GeneralReparametrizableBlock.reparametrize()fusesbranchesinto a newfused_branch, soload_checkpoint()after reparameterization will skip the original parameter keys and potentially leave quantization initialized from older weights. Apply requested weights first, then reparametrize the fused checkpoint.🐛 Proposed fix
- model.reparametrize().eval() - if weights is not None: model.load_checkpoint(weights) + model.reparametrize() + model.eval()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/core/core.py` around lines 1708 - 1721, Update _build_quant_model to load the requested checkpoint on the selected model before calling reparametrize(). Keep eval() after reparameterization, ensuring checkpoint keys for the original branches are restored before they are fused and the quantization model uses the requested weights.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@luxonis_train/core/core.py`:
- Around line 634-642: Update the _upload_export_artifact call in the YAML
export flow to pass yaml_path.name as the artifact name instead of
str(yaml_path), ensuring the remote key contains only the filename and not the
local filesystem path.
---
Outside diff comments:
In `@luxonis_train/core/core.py`:
- Around line 707-713: Update _build_modelconverter_config so the input_model
value is converted to a string before the returned configuration is passed to
YAML serialization, while preserving the existing path value and all other
configuration fields.
- Around line 1708-1721: Update _build_quant_model to load the requested
checkpoint on the selected model before calling reparametrize(). Keep eval()
after reparameterization, ensuring checkpoint keys for the original branches are
restored before they are fused and the quantization model uses the requested
weights.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: bdb54088-2d34-426f-b185-e1472f64bfe6
📒 Files selected for processing (10)
luxonis_train/attached_modules/base_attached_module.pyluxonis_train/attached_modules/visualizers/utils.pyluxonis_train/config/config.pyluxonis_train/core/core.pyluxonis_train/lightning/luxonis_lightning.pyluxonis_train/nodes/base_node.pyluxonis_train/upgrade.pyluxonis_train/utils/annotation.pyluxonis_train/utils/boundingbox.pytests/integration/test_predefined_models.py
🚧 Files skipped from review as they are similar to previous changes (8)
- luxonis_train/attached_modules/base_attached_module.py
- luxonis_train/attached_modules/visualizers/utils.py
- tests/integration/test_predefined_models.py
- luxonis_train/upgrade.py
- luxonis_train/nodes/base_node.py
- luxonis_train/lightning/luxonis_lightning.py
- luxonis_train/utils/boundingbox.py
- luxonis_train/config/config.py
kozlov721
marked this pull request as draft
July 29, 2026 14:11
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Specification
complexipytopre-commitDependencies & Potential Impact
None / not applicable
Deployment Plan
None / not applicable
Testing & Validation
None / not applicable
AI Usage
Assisted-by: AGENT_NAME:MODEL_VERSION [TOOL1] [TOOL2]
Submitted code was reviewed by a human: YES/NO
The author is taking the responsibility for the contribution: YES/NO
Summary by CodeRabbit