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 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:
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 subnetUflowModel
: 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 modelUFlowLoss
: Loss function for trainingAnomalyMapGenerator
: 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:
Start with input nodes matching feature extractor outputs
- 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
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 stagesAllInOneBlock
: 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 implementationAffineCouplingSubnet
: 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:
- 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, returnsInferenceBatch
with anomaly scores and map.
- Return type:
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 architectureUFlowLoss
: Loss function for trainingAnomalyMapGenerator
: 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 toTrue
.post_processor (PostProcessor | bool, optional) – Post-processor for model outputs. If
True
, uses default post-processor. Defaults toTrue
.evaluator (Evaluator | bool, optional) – Evaluator for model performance. If
True
, uses default evaluator. Defaults toTrue
.visualizer (Visualizer | bool, optional) – Visualizer for model outputs. If
True
, uses default visualizer. Defaults toTrue
.
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 architectureUFlowLoss
: Loss function for trainingAnomalyMapGenerator
: Anomaly map generation from features
- configure_optimizers()#
Configure optimizers and learning rate schedulers.
- 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:
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
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:
Computing per-scale likelihoods from latent variables
Upscaling likelihoods to original image size
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 mapscompute_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:
- 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
- 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
compute_anomaly_map()
: Main method for likelihood-based mapscompute_anomaly_mask()
: Optional method for NFA-based segmentation
- 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:
- 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: