Geometry of the Visual Cortex with Applications to Image Inpainting and Enhancement

May 15, 2026 · View on GitHub

Paper: Francesco Ballerin and Erlend Grong, Geometry of the visual cortex with applications to image inpainting and enhancement, Graphical Models, Volume 137, 2025, 101239, ISSN 1524-0703, Elsevier. ScienceDirect | arXiv preprint

Code: github.com/ballerin/v1diffusion

Introduction

This repository implements image inpainting and enhancement algorithms grounded in a sub-Riemannian model of the primary visual cortex V1. The key insight, originally formalised by Petitot and later exploited by Citti, Sarti, and Boscain et al., is that V1 can be modelled geometrically as the rototranslation group SE(2), where neurons are indexed by their spatial position (x, y) and their preferred orientation θ. Inpainting then becomes a matter of running a hypoelliptic diffusion equation on a 3D lifted representation of the image.

The algorithms in this repository innovate on previous implementations in three ways:

  1. Gaussian lift — Instead of lifting the image to a Dirac mass concentrated on the level-curve manifold (which causes rapid fading), we spread neural activity in the θ-fibre as a Gaussian centred at the level-curve orientation. This prevents noise and fading even for long diffusion times.

  2. WaxOn–WaxOff — By alternating hypoelliptic diffusion along level curves (WaxOn, operator X₁²+βX₂²) with reverse diffusion transversal to level curves (WaxOff, operator X₃²+βX₂²), we achieve inpainting while actively preserving sharpness.

  3. SE(2) unsharp filter — A direct analogue of the classical 2-D unsharp mask, defined entirely in SE(2), giving a curvature-sensitive sharpening tool that is particularly effective on retinal scans and other images with curvilinear structures.

OriginalCorrupted1× WaxOn–WaxOff3× WaxOn–WaxOff

Sub-Riemannian geometry of SE(2)

The special Euclidean group SE(2) is the 3-dimensional Lie group of planar roto-translations, with coordinates (x, y, θ). Its left-invariant vector fields are

X1=cos⁡θ ∂x+sin⁡θ ∂y,X2=∂θ,X3=−sin⁡θ ∂x+cos⁡θ ∂y.X_1 = \cos\theta\,\partial_x + \sin\theta\,\partial_y, \qquad X_2 = \partial_\theta, \qquad X_3 = -\sin\theta\,\partial_x + \cos\theta\,\partial_y.

The pair (H, g_β) with H = span{X₁, X₂} and g_β making (X₁, β^{-1/2}X₂) orthonormal defines a sub-Riemannian structure. Because [X₁, X₂] = −X₃, the distribution H is bracket-generating, which (by Chow–Rashevskii) means any two points can be connected by a horizontal path. The corresponding sub-Laplacian

$$\Delta_\beta = X_121^{2} + \beta X_$2^{2}$$$

is hypoelliptic (not elliptic) but still admits a smooth, strictly positive heat kernel. This is the core diffusion operator used for inpainting.

A second sub-Riemannian structure H̃ = span{X₃, X₂} gives the transversal operator

\tilde\Delta_\beta = X_$3^{2}$ + \beta X_$2^{2}$,

whose diffusion blurs across level curves. WaxOff is the reverse of this diffusion.

Lifting an image to SE(2)

ImageOrientation fieldLifted image
circles r2circles thetacircles se2

Gaussian lift (eq. 4.1 of the paper): Each pixel (x, y) is mapped to a full θ-slice in SE(2) via a Gaussian centred at the level-curve angle:

I~(x,y,θ)=I(x,y)⋅exp⁡ ⁣(−∣∇I∣2−(X3I)22σ2∣∇I∣2).\tilde I(x,y,\theta) = I(x,y)\cdot\exp\!\left(-\frac{|\nabla I|^2 - (X_3 I)^2}{2\sigma^2|\nabla I|^2}\right).

The lift is π-periodic in θ (so it lives on PTR²), and σ controls how tightly angular activity is concentrated. For σ → 0 the classical Dirac-mass lift is recovered.

Effect of σ on a corrupted eye image (all at T=10, β=0.25):

σ = 0.8 (fading)σ = 100 (stable)
sigma 0.8sigma 100

Comparison: X₁ diffusion, X₃ diffusion, WaxOn–WaxOff

OriginalAlong X₁ (WaxOn)Along X₃ (WaxOff)WaxOn–WaxOff

WaxOn diffuses along level curves to complete gaps. WaxOff concentrates across level curves to sharpen. Their alternation achieves both.

SE(2) unsharp filter

The classical 2-D unsharp filter is I + C·(I − I⊛G). The SE(2) analogue applies forward transversal diffusion (X₃²+βX₂²) to the lifted image and subtracts the result:

I~+C ⁣(I~−eTΔ~βI~).\tilde I + C\!\left(\tilde I - e^{T\tilde\Delta_\beta}\tilde I\right).

This is curvature-sensitive: edges that curve in the plane are enhanced because X₃ follows their normal direction in 3D.

Original retinal scanClassical R² unsharpSE(2) unsharp
retinalr2se2

AHE inpainting (known-mask setting)

When the corruption mask is known the AHE algorithm of Boscain et al. (2018) applies: BFS-fill → strong diffusion → advanced BFS blend → weak diffusion. This repository extends AHE by replacing the plain lift with the Gaussian lift and replacing plain diffusion passes with WaxOn–WaxOff, giving sharper reconstructions with higher contrast.

Masked inputOur AHE result

Broken circle completion (classic algorithm, β = 0.25):

OriginalRestored
circle origcircle 0.25
#CitationContribution
[1]Citti & Sarti (2006). A cortical based model of perceptual completion in the roto-translation space. JMIV 24(3).Foundational SE(2) perceptual completion model; mean-curvature flow for curve restoration in roto-translation space.
[2]Boscain, Duplaix, Gauthier, Rossi (2012). Anthropomorphic image reconstruction via hypoelliptic diffusion. SIAM J. Control Optim. 50(3).Introduced hypoelliptic diffusion (Δ_β = X₁²+βX₂²) as a blind inpainting operator.
[3]Boscain, Gauthier, Prandi, Remizov (2014). Image reconstruction via non-isotropic diffusion in Riemannian and sub-Riemannian spaces. CDC 2014.Semi-discrete implementation; variant where corruption position is known.
[4]Boscain, Chertoviskih, Gauthier, Prandi, Remizov (2018). Highly corrupted image inpainting through hypoelliptic diffusion. J. Math. Imaging Vision.AHE algorithm: BFS fill + strong/weak varying-coefficient diffusion + advanced averaging.
[5]Franken & Duits (2009); Duits et al. (2010).Enhancement and sharpening in SE(2) orientation scores.
[6]Bekkers et al. (2018). Roto-translation covariant CNNs. MICCAI 2018.SE(2) group convolution layer for geometric deep learning on retinal/histology images.

Repository Structure

v1diffusion/
├── operators.py       # JAX-JIT directional derivative operators on 3D lifted arrays
├── transforms.py      # Image lift (Gaussian, gradient, Gabor) and projection operators
├── evolution.py       # Time-integration loops for diffusion and concentration
├── processing.py      # High-level AHE inpainting pipeline
├── examples.py        # Synthetic test image generators
├── utils.py           # BFS mask fill, imshow, animation saving, angle utilities
├── metrics.py         # MSE metric
├── requirements.txt   # Python dependencies
├── pyproject.toml     # Project metadata and ruff linter configuration
├── images/            # Sample images used in notebooks and README
│   └── readme/        # Figures from the paper (preserved here)
└── jupyter_notebooks/ # Reproducible demo notebooks (see below)

operators.py

Implements the six fundamental operators on 3D arrays (H, W, N_theta) using JAX for JIT compilation:

SymbolFunctionMeaning
X₁dX1Derivative along level curves
X₂dX2Derivative w.r.t. θ (angular)
X₃dX3Derivative transversal to level curves
X₁²+βX₂²hypoelliptic_heat_operatorSub-Laplacian for WaxOn
X₃²+βX₂²orthogonal_operatorTransversal operator for WaxOff
a·X₁²+b·X₂²vc_hypoelliptic_heat_operatorVarying-coefficient operator for AHE

Both SE(2) (model=0, θ∈[0,2π)) and PTR² (model=1, θ∈[0,π)) are supported.

transforms.py

FunctionDescription
lift_gaussian(I, sigma)Gaussian lift (eq. 4.1); always π-periodic
lift_gradient(I)Dirac-mass lift via gradient direction
lift_gradient_rossi(I)Gradient lift with Rossi smoothing
lift_gabor(I, sigma)Gabor-filter lift
project_max(f)Max projection Π_max
project_integral(f)Integral projection over θ-fibres
project_gaussian(f, sigma)Gaussian inverse projection Π_σ

evolution.py

FunctionDescription
parallel_diffusionWaxOn: forward X₁²+βX₂² diffusion
orthogonal_diffusionX₃²+βX₂² forward diffusion (blur across level curves)
parallel_concentrationWaxOff via reverse Δ_β
orthogonal_concentrationWaxOff via reverse X₃²+βX₂²
v1_unsharp_filterSE(2) unsharp: u + C·(u − blur(u))
evolve_vc_hypoellipticVarying-coefficient AHE diffusion step

processing.py

AHE(image, mask, ...) — full AHE pipeline (BFS fill → strong vc-diffusion → advanced BFS blend → weak vc-diffusion → projection).

utils.py

fill_mask (BFS), imshow, ang_d (angular distance).

Jupyter Notebooks

All notebooks are in jupyter_notebooks/ and are self-contained with markdown explanations.

NotebookWhat it demonstrates
gaussian_lift.ipynbEffect of σ in the Gaussian lift; σ-sweep comparison on a corrupted eye image
waxon_waxoff_basic_example.ipynbWaxOn-only vs WaxOn–WaxOff on a synthetic horizontal-lines image with a vertical gap
waxon_waxoff_eye_example.ipynbIterated WaxOn–WaxOff on the AI-generated eye image; full parameter sweep
waxon_waxoff_unsharp_filter.ipynbTight-σ WaxOn variant combined with SE(2) unsharp for sharper inpainting
SE2_unsharp.ipynbSE(2) unsharp filter on a retinal scan; comparison with classical R² unsharp
AHE.ipynbSingle-pass AHE on a portrait image with 95%-random corruption

Installation

Recommended: conda environment

conda create -n v1diffusion python=3.10
conda activate v1diffusion
pip install -r requirements.txt

For GPU support replace jax[cpu] in requirements.txt with the appropriate JAX CUDA variant from jax.readthedocs.io.

Linting (developers):

pip install ruff
ruff check .

Ruff is configured in pyproject.toml (line-length 100, rules E/W/F/I/UP/B/A/NPY).

Citation

If you use this code or build on the methods, please cite:

@article{ballerin2025geometry,
  title     = {Geometry of the Visual Cortex with Applications to Image Inpainting and Enhancement},
  author    = {Ballerin, Francesco and Grong, Erlend},
  journal   = {Graphical Models},
  volume    = {136},
  year      = {2025},
  publisher = {Elsevier},
  url       = {https://www.sciencedirect.com/science/article/pii/S1524070324000274}
}

Preprint:

@misc{ballerin2023geometry,
  title        = {Geometry of the Visual Cortex with Applications to Image Inpainting and Enhancement},
  author       = {Ballerin, Francesco and Grong, Erlend},
  year         = {2023},
  eprint       = {2308.07652},
  archivePrefix = {arXiv},
  primaryClass = {math.AP},
  url          = {https://arxiv.org/abs/2308.07652}
}

References

  • [1] Citti, G. and Sarti, A. (2006). A cortical based model of perceptual completion in the roto-translation space. Journal of Mathematical Imaging and Vision, 24(3):307–326.
  • [2] Boscain, U., Duplaix, J., Gauthier, J. P., and Rossi, F. (2012). Anthropomorphic image reconstruction via hypoelliptic diffusion. SIAM Journal on Control and Optimization, 50(3):1309–1336.
  • [3] Boscain, U., Gauthier, J. P., Prandi, D., and Remizov, A. (2014). Image reconstruction via non-isotropic diffusion in Riemannian and sub-Riemannian spaces. In 53rd IEEE Conference on Decision and Control, pages 4958–4963.
  • [4] Boscain, U., Chertoviskih, R., Gauthier, J. P., Prandi, D., and Remizov, A. (2018). Highly corrupted image inpainting through hypoelliptic diffusion. Journal of Mathematical Imaging and Vision.
  • [5] Franken, E. and Duits, R. (2009). Crossing-preserving coherence-enhancing diffusion on invertible orientation scores. IJCV, 85(3):253–278.
  • [6] Bekkers, E. J. et al. (2018). Roto-translation covariant convolutional networks for medical image analysis. MICCAI 2018.

License

MIT