U-Flow#

This is the implementation of the U-Flow paper.

Model Type: Segmentation

Description#

U-Flow is a U-Shaped normalizing flow-based probability distribution estimator. The method consists of three phases. (1) Multi-scale feature extraction: a rich multi-scale representation is obtained with MSCaiT, by combining pre-trained image Transformers acting at different image scales. It can also be used any other feature extractor, such as ResNet. (2) U-shaped Normalizing Flow: by adapting the widely used U-like architecture to NFs, a fully invertible architecture is designed. This architecture is capable of merging the information from different scales while ensuring independence both intra- and inter-scales. To make it fully invertible, split and invertible up-sampling operations are used. (3) Anomaly score and segmentation computation: besides generating the anomaly map based on the likelihood of test data, we also propose to adapt the a contrario framework to obtain an automatic threshold by controlling the allowed number of false alarms.

Architecture#

U-Flow Architecture

U-Flow PyTorch Implementation.

This module provides the PyTorch implementation of the U-Flow model for anomaly detection. U-Flow combines normalizing flows with a U-Net style architecture to learn the distribution of normal images and detect anomalies.

The model consists of several key components:
  • Feature extraction using a pre-trained backbone

  • Normalizing flow blocks arranged in a U-Net structure

  • Anomaly map generation for localization

The implementation includes classes for:
  • Affine coupling subnet construction

  • Main U-Flow model architecture

  • Anomaly map generation

class anomalib.models.image.uflow.torch_model.AffineCouplingSubnet(kernel_size, subnet_channels_ratio)#

Bases: object

Class for building the Affine Coupling subnet.

This class creates a subnet used within the affine coupling layers of the normalizing flow. The subnet is passed as an argument to the AllInOneBlock module and determines how features are transformed within the coupling layer.

Parameters:
  • kernel_size (int) – Size of convolutional kernels used in subnet layers.

  • subnet_channels_ratio (float) – Ratio determining the number of intermediate channels in the subnet relative to input channels.

Example

>>> subnet = AffineCouplingSubnet(kernel_size=3, subnet_channels_ratio=1.0)
>>> layer = subnet(in_channels=64, out_channels=128)
>>> layer
Sequential(
  (0): Conv2d(64, 64, kernel_size=(3, 3), padding=same)
  (1): ReLU()
  (2): Conv2d(64, 128, kernel_size=(3, 3), padding=same)
)

See also

  • AllInOneBlock: Flow block using this subnet

  • UflowModel: Main model incorporating these subnets

class anomalib.models.image.uflow.torch_model.UflowModel(input_size=(448, 448), flow_steps=4, backbone='mcait', affine_clamp=2.0, affine_subnet_channels_ratio=1.0, permute_soft=False)#

Bases: Module

PyTorch implementation of the U-Flow model architecture.

This class implements the U-Flow model for anomaly detection. The model consists of:

  • A U-shaped normalizing flow architecture for density estimation

  • Multi-scale feature extraction using pre-trained backbones

  • Unsupervised threshold estimation based on the learned density

  • Anomaly detection by comparing likelihoods to the threshold

Parameters:
  • input_size (tuple[int, int]) – Input image dimensions as (height, width). Defaults to (448, 448).

  • flow_steps (int) – Number of normalizing flow steps in each flow stage. Defaults to 4.

  • backbone (str) – Name of the backbone feature extractor. Must be one of ["mcait", "resnet18", "wide_resnet50_2"]. Defaults to "mcait".

  • affine_clamp (float) – Clamping value for affine coupling layers. Defaults to 2.0.

  • affine_subnet_channels_ratio (float) – Channel ratio for affine coupling subnet. Defaults to 1.0.

  • permute_soft (bool) – Whether to use soft permutation. Defaults to False.

Example

>>> import torch
>>> from anomalib.models.image.uflow.torch_model import UflowModel
>>> model = UflowModel(
...     input_size=(256, 256),
...     backbone="resnet18"
... )
>>> image = torch.randn(1, 3, 256, 256)
>>> output = model(image)  # Returns anomaly map during inference

See also

  • Uflow: Lightning implementation using this model

  • UFlowLoss: Loss function for training

  • AnomalyMapGenerator: Anomaly map generation from features

build_flow(flow_steps)#

Build the U-shaped normalizing flow architecture.

The flow is built in a U-shaped structure, processing features from coarse to fine scales:

  1. Start with input nodes matching feature extractor outputs

  2. For each scale (coarse to fine):
    • Pass through flow stage (sequence of coupling layers)

    • Split output into two parts

    • Send one part to output

    • Upsample other part and concatenate with next scale

  3. Build final flow graph combining all nodes

Parameters:

flow_steps (int) – Number of coupling layers in each flow stage.

Returns:

Constructed normalizing flow graph.

Return type:

ff.GraphINN

See also

  • build_flow_stage(): Builds individual flow stages

  • AllInOneBlock: Individual coupling layer blocks

build_flow_stage(in_node, flow_steps, condition_node=None)#

Build a single flow stage consisting of multiple coupling layers.

Each flow stage is a sequence of flow_steps Glow-style coupling blocks (AllInOneBlock). The blocks alternate between 3x3 and 1x1 convolutions in their coupling subnets.

Parameters:
  • in_node (ff.Node) – Input node to the flow stage.

  • flow_steps (int) – Number of coupling layers to use.

  • condition_node (ff.Node, optional) – Optional conditioning node. Defaults to None.

Returns:

List of constructed coupling layer nodes.

Return type:

list[ff.Node]

See also

  • AllInOneBlock: Individual coupling layer implementation

  • AffineCouplingSubnet: Subnet used in coupling layers

encode(features)#

Encode input features to latent space using normalizing flow.

Parameters:

features (torch.Tensor) – Input features from feature extractor.

Returns:

Tuple containing:
  • Latent variables from flow transformation

  • Log-Jacobian determinant of the transformation

Return type:

tuple[torch.Tensor, torch.Tensor]

forward(image)#

Process input image through the model.

During training, returns latent variables and log-Jacobian determinant. During inference, returns anomaly scores and anomaly map.

Parameters:

image (torch.Tensor) – Input image tensor of shape (batch_size, channels, height, width).

Returns:

During training, returns tuple of

(latent_vars, log_jacobian). During inference, returns InferenceBatch with anomaly scores and map.

Return type:

torch.Tensor | InferenceBatch

U-Flow: A U-shaped Normalizing Flow for Anomaly Detection with Unsupervised Threshold.

This module implements the U-Flow model for anomaly detection as described in

<https://arxiv.org/pdf/2211.12353.pdf>`_. The model consists

of:

  • A U-shaped normalizing flow architecture for density estimation

  • Multi-scale feature extraction using pre-trained backbones

  • Unsupervised threshold estimation based on the learned density

  • Anomaly detection by comparing likelihoods to the threshold

Example

>>> from anomalib.models.image import Uflow
>>> from anomalib.engine import Engine
>>> from anomalib.data import MVTec
>>> datamodule = MVTec()
>>> model = Uflow()
>>> engine = Engine(model=model, datamodule=datamodule)
>>> engine.fit()  
>>> predictions = engine.predict()  

See also

  • UflowModel: PyTorch implementation of the model architecture

  • UFlowLoss: Loss function for training

  • AnomalyMapGenerator: Anomaly map generation from features

class anomalib.models.image.uflow.lightning_model.Uflow(backbone='mcait', flow_steps=4, affine_clamp=2.0, affine_subnet_channels_ratio=1.0, permute_soft=False, pre_processor=True, post_processor=True, evaluator=True, visualizer=True)#

Bases: AnomalibModule

Lightning implementation of the U-Flow model.

This class implements the U-Flow model for anomaly detection as described in Rudolph et al., 2022. The model consists of:

  • A U-shaped normalizing flow architecture for density estimation

  • Multi-scale feature extraction using pre-trained backbones

  • Unsupervised threshold estimation based on the learned density

  • Anomaly detection by comparing likelihoods to the threshold

Parameters:
  • backbone (str, optional) – Name of the backbone feature extractor. Must be one of ["mcait", "resnet18", "wide_resnet50_2"]. Defaults to "mcait".

  • flow_steps (int, optional) – Number of normalizing flow steps. Defaults to 4.

  • affine_clamp (float, optional) – Clamping value for affine coupling layers. Defaults to 2.0.

  • affine_subnet_channels_ratio (float, optional) – Channel ratio for affine coupling subnet. Defaults to 1.0.

  • permute_soft (bool, optional) – Whether to use soft permutation. Defaults to False.

  • pre_processor (PreProcessor | bool, optional) – Pre-processor for input data. If True, uses default pre-processor. Defaults to True.

  • post_processor (PostProcessor | bool, optional) – Post-processor for model outputs. If True, uses default post-processor. Defaults to True.

  • evaluator (Evaluator | bool, optional) – Evaluator for model performance. If True, uses default evaluator. Defaults to True.

  • visualizer (Visualizer | bool, optional) – Visualizer for model outputs. If True, uses default visualizer. Defaults to True.

Example

>>> from anomalib.models.image import Uflow
>>> from anomalib.engine import Engine
>>> from anomalib.data import MVTec
>>> datamodule = MVTec()
>>> model = Uflow(backbone="resnet18")
>>> engine = Engine(model=model, datamodule=datamodule)
>>> engine.fit()  
>>> predictions = engine.predict()  
Raises:

ValueError – If input_size is not provided during initialization.

See also

  • UflowModel: PyTorch implementation of the model architecture

  • UFlowLoss: Loss function for training

  • AnomalyMapGenerator: Anomaly map generation from features

configure_optimizers()#

Configure optimizers and learning rate schedulers.

Returns:

Tuple containing:
  • List of optimizers (Adam with initial lr=1e-3)

  • List of schedulers (LinearLR reducing to 0.4 over 25000 steps)

Return type:

tuple[list[LightningOptimizer], list[LRScheduler]]

classmethod configure_pre_processor(image_size=None)#

Configure default pre-processor for U-Flow model.

Parameters:

image_size (tuple[int, int] | None, optional) – Input image size. Not used as U-Flow has fixed input size. Defaults to None.

Returns:

Default pre-processor with resizing and normalization.

Return type:

PreProcessor

Note

The input image size is fixed to 448x448 for U-Flow regardless of the provided image_size.

property learning_type: LearningType#

Get the learning type of the model.

Returns:

Learning type (ONE_CLASS for U-Flow)

Return type:

LearningType

property trainer_arguments: dict[str, Any]#

Get trainer arguments for U-Flow.

Returns:

Dictionary containing trainer arguments

Return type:

dict[str, Any]

training_step(batch, *args, **kwargs)#

Perform a training step.

Parameters:
  • batch (Batch) – Input batch containing images

  • *args – Variable length argument list

  • **kwargs – Arbitrary keyword arguments

Returns:

Dictionary containing the loss value

Return type:

STEP_OUTPUT

validation_step(batch, *args, **kwargs)#

Perform a validation step.

Parameters:
  • batch (Batch) – Input batch containing images

  • *args – Variable length argument list

  • **kwargs – Arbitrary keyword arguments

Returns:

Batch updated with model predictions

Return type:

STEP_OUTPUT

Anomaly map computation for U-Flow model.

This module implements functionality to generate anomaly heatmaps from the latent variables produced by a U-Flow model. The anomaly maps are generated by:

  1. Computing per-scale likelihoods from latent variables

  2. Upscaling likelihoods to original image size

  3. Combining multiple scale likelihoods

Example

>>> from anomalib.models.image.uflow.anomaly_map import AnomalyMapGenerator
>>> generator = AnomalyMapGenerator(input_size=(256, 256))
>>> latent_vars = [torch.randn(1, 64, 32, 32), torch.randn(1, 128, 16, 16)]
>>> anomaly_map = generator(latent_vars)

See also

  • AnomalyMapGenerator: Main class for generating anomaly maps

  • compute_anomaly_map(): Function to generate anomaly maps from latents

class anomalib.models.image.uflow.anomaly_map.AnomalyMapGenerator(input_size)#

Bases: Module

Generate anomaly heatmaps and segmentation masks from U-Flow latent variables.

This class implements functionality to generate anomaly maps by analyzing the latent variables produced by a U-Flow model. The anomaly maps can be generated in two ways:

  1. Using likelihood-based scoring (default method):
    • Computes per-scale likelihoods from latent variables

    • Upscales likelihoods to original image size

    • Combines multiple scale likelihoods via averaging

  2. Using NFA-based segmentation (optional method):
    • Applies binomial testing on local windows

    • Computes Number of False Alarms (NFA) statistics

    • Generates binary segmentation masks

Parameters:

input_size (ListConfig | tuple) – Size of input images as (height, width)

Example

>>> from anomalib.models.image.uflow.anomaly_map import AnomalyMapGenerator
>>> generator = AnomalyMapGenerator(input_size=(256, 256))
>>> latents = [torch.randn(1, 64, 32, 32), torch.randn(1, 128, 16, 16)]
>>> anomaly_map = generator(latents)
>>> anomaly_map.shape
torch.Size([1, 1, 256, 256])

See also

static binomial_test(z, window_size, probability_thr, high_precision=False)#

Apply binomial test to validate/reject normality hypothesis.

For each pixel, tests the null hypothesis that the pixel and its local neighborhood are normal against the alternative that they are anomalous.

The test: 1. Counts anomalous pixels in local window using chi-square threshold 2. Compares observed count to expected count under null hypothesis 3. Returns log probability of observing such extreme counts

Parameters:
  • z (torch.Tensor) – Latent tensor of shape (batch_size, channels, height, width)

  • window_size (int) – Size of local window for counting

  • probability_thr (float) – Probability threshold for chi-square test

  • high_precision (bool, optional) – Whether to use high precision computation. Defaults to False.

Returns:

Log probability tensor of shape ``(batch_size, 1,

height, width)``

Return type:

torch.Tensor

compute_anomaly_map(latent_variables)#

Generate likelihood-based anomaly map from latent variables.

The method: 1. Computes per-scale likelihoods from latent variables 2. Upscales each likelihood map to input image size 3. Combines scale likelihoods via averaging

Parameters:

latent_variables (list[Tensor]) – List of latent tensors from U-Flow model, each with shape (batch_size, channels, height, width)

Returns:

Anomaly heatmap of shape (batch_size, 1, height, width)

Return type:

Tensor

compute_anomaly_mask(z, window_size=7, binomial_probability_thr=0.5, high_precision=False)#

Generate NFA-based anomaly segmentation mask from latent variables.

This optional method implements the Number of False Alarms (NFA) approach from the U-Flow paper. It is slower than the default likelihood method but provides unsupervised binary segmentation.

The method: 1. Applies binomial testing on local windows around each pixel 2. Computes NFA statistics based on concentration of candidate pixels 3. Generates binary segmentation mask

Parameters:
  • z (list[torch.Tensor]) – List of latent tensors from U-Flow model

  • window_size (int, optional) – Size of local window for binomial test. Defaults to 7.

  • binomial_probability_thr (float, optional) – Probability threshold for binomial test. Defaults to 0.5.

  • high_precision (bool, optional) – Whether to use high precision NFA computation. Slower but more accurate. Defaults to False.

Returns:

Binary anomaly mask of shape ``(batch_size, 1, height,

width)``

Return type:

torch.Tensor

forward(latent_variables)#

Generate anomaly map from latent variables.

Parameters:

latent_variables (list[Tensor]) – List of latent tensors from U-Flow model

Returns:

Anomaly heatmap of shape (batch_size, 1, height, width)

Return type:

Tensor