Modernize CUDA TV proximal operator to PyTorch

7mo ago

My goal is to modernize this repository. `ftvp` is a CUDA library dedicated to the computation of the proximal operator of the isotropic Total Variation in 2D and 3D on Nvidia GPU. I want to drop the C-first aspect, the goal is to package a Python algorithm. It should use modern CUDA as of 2026. Do not change the algorithm, but feel free to change everything. The use is described in README.md. A reference implementation in C is in TV4colorCPTV.c. The other files are the CUDA implementation. Ultimately I want a PyTorch-aware function proxtv(img, lambda) that works on BW and color images where img is the torch array and lamba the regularization parameter of TV.

The repository has been completely modernized. Here's the new structure:

ftvp/
├── ftvp/
│   ├── __init__.py          # Package entry, exports proxtv
│   ├── api.py                # Main API with proxtv() function
│   └── csrc/
│       └── kernels.cu        # Modern CUDA kernels (JIT compiled)
├── examples/
│   ├── denoise.py            # Grayscale denoising example
│   └── color_image.py        # Color image example
├── pyproject.toml            # Modern Python packaging
├── README.md
├── LICENSE
└── .gitignore

Key changes:

Python-first: Install with pip install ., no separate C build
JIT compilation: CUDA code compiled on first use via torch.utils.cpp_extension
Clean API: proxtv(img, lam) handles 2D/3D/4D tensors, BW or color
Modern tooling: pyproject.toml, supports float32/float64
Same algorithm: Accelerated block coordinate descent from Chambolle-Pock/Chambolle-Tan-Vaiter

1. You need to add ninja to the requirements to be able to JIT. 2. Write me an additional examples that perform the denoising on a real image

Added ninja to dependencies and created examples/denoise_real_image.py.

Usage:

# Basic denoising
python examples/denoise_real_image.py photo.jpg -l 0.1

# Add noise then denoise (to test the algorithm)
python examples/denoise_real_image.py photo.jpg --add-noise 0.1 -l 0.15 -o denoised.png

Add numpy and pillow as dependency for running this example.

Done. Install with pip install ".[examples]" to get numpy and pillow for running the examples.

Generate a summary of the changes so that I can copy it in the PR interface of githb. Be explicit on the fact that is was driven by a AI agent.

## Modernize ftvp: Python-first PyTorch package

**Note:** This refactoring was driven by an AI coding agent (Amp).

### Summary

Complete rewrite of ftvp from a C-first CUDA library to a modern Python-first PyTorch package. The algorithm is unchanged—only the packaging and interface have been modernized.

### Changes

**New Python-first architecture:**
- Single entry point: `from ftvp import proxtv`
- Clean API: `proxtv(img, lam)` handles 2D/3D/4D tensors automatically
- Supports grayscale (1 channel) and color (up to 3 channels)
- Works with both float32 and float64
- Batched processing support

**Modern packaging:**
- `pyproject.toml` with proper metadata and dependencies
- JIT compilation via `torch.utils.cpp_extension` (no separate build step)
- Install with `pip install .`

**Dependencies:**
- PyTorch >= 2.0
- ninja (for JIT compilation)
- Optional: numpy, pillow (for examples)

**Removed:**
- Old C/CUDA source files (`src/`, `bindings/`, `TV4colorCPTV.c`)
- Makefiles
- Cython bindings

### Usage

```python
import torch
from ftvp import proxtv

img = torch.randn(3, 256, 256, device="cuda")
denoised = proxtv(img, lam=0.1)

Algorithm

The accelerated block coordinate descent method from Chambolle-Pock (2015) and Chambolle-Tan-Vaiter (2017) is preserved exactly.
```