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

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.

Original	Corrupted	1× WaxOn–WaxOff	3× WaxOn–WaxOff

---

Background and Related Work

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

$$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_1^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)

Image	Orientation field	Lifted image

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:

$$\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)

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

Original	Along 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:

$$\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 scan	Classical R² unsharp	SE(2) unsharp

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 input	Our AHE result

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

Original	Restored

---

Related Papers

#	Citation	Contribution
[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:

Symbol	Function	Meaning
X₁	dX1	Derivative along level curves
X₂	dX2	Derivative w.r.t. θ (angular)
X₃	dX3	Derivative transversal to level curves
X₁²+βX₂²	hypoelliptic_heat_operator	Sub-Laplacian for WaxOn
X₃²+βX₂²	orthogonal_operator	Transversal operator for WaxOff
a·X₁²+b·X₂²	vc_hypoelliptic_heat_operator	Varying-coefficient operator for AHE

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

transforms.py

Function	Description
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

Function	Description
parallel_diffusion	WaxOn: forward X₁²+βX₂² diffusion
orthogonal_diffusion	X₃²+βX₂² forward diffusion (blur across level curves)
parallel_concentration	WaxOff via reverse Δ_β
orthogonal_concentration	WaxOff via reverse X₃²+βX₂²
v1_unsharp_filter	SE(2) unsharp: u + C·(u − blur(u))
evolve_vc_hypoelliptic	Varying-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.

Notebook	What it demonstrates
gaussian_lift.ipynb	Effect of σ in the Gaussian lift; σ-sweep comparison on a corrupted eye image
waxon_waxoff_basic_example.ipynb	WaxOn-only vs WaxOn–WaxOff on a synthetic horizontal-lines image with a vertical gap
waxon_waxoff_eye_example.ipynb	Iterated WaxOn–WaxOff on the AI-generated eye image; full parameter sweep
waxon_waxoff_unsharp_filter.ipynb	Tight-σ WaxOn variant combined with SE(2) unsharp for sharper inpainting
SE2_unsharp.ipynb	SE(2) unsharp filter on a retinal scan; comparison with classical R² unsharp
AHE.ipynb	Single-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