feat(dinomaly): Add context-aware recentering option from dinomaly2 by rajeshgangireddy · Pull Request #3435 · open-edge-platform/anomalib

rajeshgangireddy · 2026-03-19T08:08:01Z

📝 Description

DAinomaly 2 was released a few months ago. Dinomaly2 extends Dinomaly with context-Aware Recentering. The other concepts like foundation Transformers, Noisy Bottleneck, Unfocused Linear Attention, Loose Reconstruction, are identical to what's already implemented.

Context-Aware Recentering subtracts the CLS (class) token from each patch feature before reconstruction:
This conditions feature reconstruction on class-specific context, preventing the decoder from reconstructing patterns that are "normal" in one object class but anomalous in another.

The code is not public, but for single class UAD scenario, the changes are simple (at least I assume)

In this PR, added an extra 'use_context_recentering' that toggles between dinomaly2 (if true) and dinomaly implementations.

Addtionaly added a logic to show the right number of epochs when max_steps is passed (Dinomaly paper uses 5000 steps )

There could be slighly improvoement (~1 to 2%) in pixel level scores.
Example numbers for Bottle category MvTecAD with max_steps = 5000

Metric	Dinomaly	Dinomaly2	Delta
image_AUROC	1	1	0
image_F1Score	1	0.9839	-0.0161
pixel_AUROC	0.9887	0.9897	0.001
pixel_F1Score	0.802	0.8204	0.0184
train_time_s	5882.9	5267.5	-615.4

Need to test on more categories to get a clear picture.

✨ Changes

Select what type of change your PR is:

✅ Checklist

Before you submit your pull request, please make sure you have completed the following steps:

📚 I have made the necessary updates to the documentation (if applicable).
🧪 I have written tests that support my changes and prove that my fix is effective or my feature works (if applicable).
🏷️ My PR title follows conventional commit format.

For more information about code review checklists, see the Code Review Checklist.

Copilot

Pull request overview

This PR adds an optional Dinomaly2-style “Context-Aware Recentering” mode to Dinomaly and includes a Lightning hook to improve Rich progress-bar epoch display when training is configured with max_steps.

Changes:

Added use_context_recentering flag to toggle CLS-token-based recentering vs. existing Dinomaly behavior.
Implemented patch-feature recentering (subtract CLS token) in the encoder feature pipeline, and adjusted downstream token stripping accordingly.
Added a Lightning on_train_start workaround to estimate epochs for Rich progress display when max_steps is used.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
`src/anomalib/models/image/dinomaly/torch_model.py`	Adds `use_context_recentering` option, validates incompatibility with `remove_class_token`, and applies recentering in feature extraction.
`src/anomalib/models/image/dinomaly/lightning_model.py`	Plumbs the new flag into the Lightning module and monkey-patches Rich progress-bar epoch description for step-based training.
`examples/configs/model/dinomaly.yaml`	Exposes `use_context_recentering` in the example config.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

… for accurate epoch display

rajeshgangireddy · 2026-03-25T08:24:10Z

I ran some experiements on mvtec and visa datasets and compared with dinomaly1, there aren't huge or consistent improvements. But some of the categories get some great improvments. Training time is essentially the same as well.
So, I suggest merging these changes for now and revisting the implementation once the official code is released.

% improvement when using dinomaly 2:

100*(D2-D1)/D1 
(D2: Dinomaly2, D1: Dinomaly1)

Copilot

Pull request overview

Adds an optional Dinomaly2-style “Context-Aware Recentering” mode to Dinomaly and includes a Lightning/Rich progress-bar compatibility tweak for step-based training (max_steps).

Changes:

Introduces use_context_recentering flag and enforces incompatibility with remove_class_token.
Applies CLS-token subtraction to patch tokens prior to reconstruction when the flag is enabled.
Adds a on_train_start hook to adjust RichProgressBar epoch display when Lightning uses max_epochs=-1 for step-based runs.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 6 comments.

File	Description
`src/anomalib/models/image/dinomaly/torch_model.py`	Adds `use_context_recentering` flag and modifies token handling to recenter patch features.
`src/anomalib/models/image/dinomaly/lightning_model.py`	Wires the new flag through the Lightning module and patches Rich progress-bar epoch display for `max_steps` runs.
`examples/configs/model/dinomaly.yaml`	Exposes `use_context_recentering` in the example config.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

…o make copilot happy)

Copilot

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-03-26T09:36:25Z

+        if use_context_recentering and remove_class_token:
+            msg = (
+                "use_context_recentering=True requires access "
+                "to the class token and is incompatible with remove_class_token=True"
+            )


The new use_context_recentering option adds a new feature-processing path and an incompatibility constraint with remove_class_token, but there are currently no unit tests covering DinomalyModel. Since this repo has unit tests for other image torch models, please add tests that (1) assert the incompatible flag combination raises, and (2) validate output tensor shapes/paths for use_context_recentering vs baseline to prevent regressions.

Copilot · 2026-03-26T09:36:26Z

+        """Fix Rich progress bar epoch display when using step-based training.
+
+        Lightning internally sets ``max_epochs=-1`` when only ``max_steps`` is
+        provided. The Rich progress bar then displays "Epoch X/-2" because it
+        computes ``max_epochs - 1``. This hook calculates the estimated number
+        of epochs from ``max_steps`` and ``num_training_batches`` so the
+        progress bar shows "Epoch X/M", where ``M`` is the estimated last
+        epoch index (``estimated_epochs - 1``).
+
+        Note:
+            This relies on Lightning's ``RichProgressBar._get_train_description``
+            as implemented in Lightning 2.x. If the internals change in future
+            versions, this hook may need updating.
+        """
+        if self.trainer.max_epochs is not None and self.trainer.max_epochs < 0:
+            try:
+                from lightning.pytorch.callbacks.progress.rich_progress import RichProgressBar
+            except ImportError:
+                return
+
+            progress_bar = getattr(self.trainer, "progress_bar_callback", None)
+            if isinstance(progress_bar, RichProgressBar) and hasattr(
+                progress_bar,
+                "_get_train_description",
+            ):
+                num_batches = self.trainer.num_training_batches
+                max_steps = self.trainer.max_steps
+                if (
+                    max_steps > 0
+                    and isinstance(num_batches, (int, float))
+                    and num_batches > 0
+                    and math.isfinite(num_batches)
+                ):
+                    est_max_epochs = math.ceil(max_steps / num_batches)
+                else:
+                    est_max_epochs = None
+
+                val_desc = getattr(progress_bar, "validation_description", "Validation")
+
+                class _FixedRichProgressBar(RichProgressBar):
+                    """RichProgressBar subclass with corrected epoch display for step-based training."""
+
+                    def _get_train_description(self, current_epoch: int) -> str:  # noqa: PLR6301
+                        desc = f"Epoch {current_epoch}"
+                        if est_max_epochs is not None:
+                            desc += f"/{est_max_epochs - 1}"
+                        if len(val_desc) > len(desc):
+                            desc = f"{desc:{len(val_desc)}}"
+                        return desc
+
+                progress_bar.__class__ = _FixedRichProgressBar
+


This hook monkey-patches the trainer’s RichProgressBar by reassigning progress_bar.__class__ and overriding the private method _get_train_description. This is brittle (depends on Lightning internals that can change without warning) and makes behavior harder to reason about/debug. Consider moving this to a dedicated callback in anomalib (so it’s explicit/configurable) or avoid patching private methods by deriving max_epochs earlier (e.g., when building the Trainer) instead of mutating the callback at runtime.

Suggested change

"""Fix Rich progress bar epoch display when using step-based training.

Lightning internally sets ``max_epochs=-1`` when only ``max_steps`` is

provided. The Rich progress bar then displays "Epoch X/-2" because it

computes ``max_epochs - 1``. This hook calculates the estimated number

of epochs from ``max_steps`` and ``num_training_batches`` so the

progress bar shows "Epoch X/M", where ``M`` is the estimated last

epoch index (``estimated_epochs - 1``).

Note:

This relies on Lightning's ``RichProgressBar._get_train_description``

as implemented in Lightning 2.x. If the internals change in future

versions, this hook may need updating.

"""

if self.trainer.max_epochs is not None and self.trainer.max_epochs < 0:

try:

from lightning.pytorch.callbacks.progress.rich_progress import RichProgressBar

except ImportError:

return

progress_bar = getattr(self.trainer, "progress_bar_callback", None)

if isinstance(progress_bar, RichProgressBar) and hasattr(

progress_bar,

"_get_train_description",

):

num_batches = self.trainer.num_training_batches

max_steps = self.trainer.max_steps

if (

max_steps > 0

and isinstance(num_batches, (int, float))

and num_batches > 0

and math.isfinite(num_batches)

):

est_max_epochs = math.ceil(max_steps / num_batches)

else:

est_max_epochs = None

val_desc = getattr(progress_bar, "validation_description", "Validation")

class _FixedRichProgressBar(RichProgressBar):

"""RichProgressBar subclass with corrected epoch display for step-based training."""

def _get_train_description(self, current_epoch: int) -> str: # noqa: PLR6301

desc = f"Epoch {current_epoch}"

if est_max_epochs is not None:

desc += f"/{est_max_epochs - 1}"

if len(val_desc) > len(desc):

desc = f"{desc:{len(val_desc)}}"

return desc

progress_bar.__class__ = _FixedRichProgressBar

"""Hook called at the beginning of training.

This method previously adjusted the Rich progress bar epoch display for

step-based training by monkey-patching Lightning's ``RichProgressBar``.

That behavior has been removed to avoid relying on private Lightning

internals and mutating callback classes at runtime.

"""

# Intentionally left without side effects to avoid brittle monkey-patching

# of Lightning internals while preserving the public hook API.

return None

ashwinvaidya17

Thanks for adding this. Just one minor comment

ashwinvaidya17 · 2026-03-31T07:35:48Z

        self._initialize_trainable_modules(self.trainable_modules)

+    def on_train_start(self) -> None:
+        """Fix Rich progress bar epoch display when using step-based training.


I feel this isn't limited to just this model. Maybe all models can benefit from this method if they ever use step based training. Should we move this to the base class?

…eing in dinomaly

Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 4 comments.

Copilot · 2026-04-20T12:13:42Z

 from .checkpoint import ModelCheckpoint
 from .graph import GraphLogger
 from .model_loader import LoadModelCallback
+from .rich_progress_bar import MaxStepsProgressCallback
 from .tiler_configuration import TilerConfigurationCallback
 from .timer import TimerCallback



The module docstring lists which callbacks are exported, but it isn’t updated to include MaxStepsProgressCallback even though it is now part of the public API via __all__. Please update the docstring list so it matches the actual exports.

Copilot · 2026-04-20T12:13:42Z

+"""Max-steps progress callback for step-based training.
+
+Lightning internally sets ``max_epochs=-1`` when only ``max_steps`` is provided.
+The Rich progress bar then displays "Epoch X/-2" because it computes
+``max_epochs - 1``. This callback estimates the total number of epochs from
+``max_steps`` and ``num_training_batches`` so the progress bar shows a
+meaningful "Epoch X/M" instead.


The module docstring says the progress bar will show "Epoch X/M", but the implementation appends est_max_epochs - 1 (e.g. "Epoch 0/9" for 10 epochs). Please align the docstring with the actual display convention or adjust the display logic to match the documented output.

Copilot · 2026-04-20T12:13:42Z

+            from lightning.pytorch.callbacks.progress.rich_progress import RichProgressBar
+        except ImportError:
+            return


RichProgressBar is imported from lightning.pytorch.callbacks.progress.rich_progress, which is an internal module path and more likely to change across Lightning versions. Prefer importing from the public API (e.g. lightning.pytorch.callbacks) with a fallback to the internal path if needed, to reduce breakage risk.

Suggested change

from lightning.pytorch.callbacks.progress.rich_progress import RichProgressBar

except ImportError:

return

from lightning.pytorch.callbacks import RichProgressBar

except ImportError:

try:

from lightning.pytorch.callbacks.progress.rich_progress import RichProgressBar

except ImportError:

return

Copilot · 2026-04-20T12:13:43Z

+        if use_context_recentering and remove_class_token:
+            msg = (
+                "use_context_recentering=True requires access "
+                "to the class token and is incompatible with remove_class_token=True"
+            )
+            raise ValueError(msg)
+


The new use_context_recentering behavior (including the ValueError when combined with remove_class_token=True) isn’t covered by tests. Please add unit tests that (1) assert the incompatibility check raises, and (2) verify the recentering path produces patch-only tokens and preserves the expected spatial reshape downstream.

feat(dinomaly): add context-aware recentering option from dinomaly2

1cb5ea2

Copilot AI review requested due to automatic review settings March 19, 2026 08:08

Merge branch 'main' into feature/dinomaly2

8173e90

Copilot AI reviewed Mar 19, 2026

View reviewed changes

Comment thread src/anomalib/models/image/dinomaly/lightning_model.py Outdated

Comment thread src/anomalib/models/image/dinomaly/lightning_model.py Outdated

Comment thread src/anomalib/models/image/dinomaly/lightning_model.py Outdated

fix(dinomaly): improve compatibility with Lightning's RichProgressBar…

a4e73d4

… for accurate epoch display

rajeshgangireddy marked this pull request as ready for review March 25, 2026 08:17

rajeshgangireddy requested a review from samet-akcay as a code owner March 25, 2026 08:17

Copilot AI review requested due to automatic review settings March 25, 2026 08:17

rajeshgangireddy requested a review from ashwinvaidya17 as a code owner March 25, 2026 08:17

Merge branch 'main' into feature/dinomaly2

6e1b0ba

Copilot AI reviewed Mar 25, 2026

View reviewed changes

Copilot started reviewing on behalf of rajeshgangireddy March 25, 2026 08:44 View session

rajeshgangireddy added 2 commits March 25, 2026 14:24

fix: update RichProgressBar integration for improved epoch display (t…

8d4d8b1

…o make copilot happy)

Merge branch 'main' into feature/dinomaly2

a830a4c

Copilot AI review requested due to automatic review settings March 26, 2026 09:31

Copilot started reviewing on behalf of rajeshgangireddy March 26, 2026 09:32 View session

Copilot AI reviewed Mar 26, 2026

View reviewed changes

ashwinvaidya17 reviewed Mar 31, 2026

View reviewed changes

rajeshgangireddy added 2 commits March 31, 2026 14:20

Merge branch 'main' into feature/dinomaly2

4da12b9

feat(callbacks): Moveprogressbarcall to common directory instead of b…

3b243e8

…eing in dinomaly

Copilot AI review requested due to automatic review settings April 20, 2026 12:07

Copilot started reviewing on behalf of rajeshgangireddy April 20, 2026 12:08 View session

Copilot AI reviewed Apr 20, 2026

View reviewed changes

Conversation

rajeshgangireddy commented Mar 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📝 Description

✨ Changes

✅ Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

rajeshgangireddy commented Mar 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

ashwinvaidya17 left a comment

Choose a reason for hiding this comment

Uh oh!

ashwinvaidya17 Mar 31, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Apr 20, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 20, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 20, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 20, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

rajeshgangireddy commented Mar 19, 2026 •

edited

Loading

rajeshgangireddy commented Mar 25, 2026 •

edited

Loading