jax-enhance is a jax implementation of the pytorch-enhance library. This is mostly for my own education/experimentation with jax however you may find these implementations useful.
pip install git+https://github.com/isaaccorley/jax-enhance
The following models are currently implemented:
- SRCNN from Dong et. al Image Super-Resolution Using Deep Convolutional Networks
- VDSR from Lee et al. Accurate Image Super-Resolution Using Very Deep Convolutional Networks
- ESPCN from Shi et. al Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network
- SRResNet from Ledig et. al Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network
- EDSR from Lim et. al Enhanced Deep Residual Networks for Single Image Super-Resolution
import jax
import jax.numpy as jnp
from jax.random import PRNGKey
import jax_enhance
# increase resolution by factor of 2 (e.g. 128x128 -> 256x256)
model = jax_enhance.models.SRResNet(scale_factor=2, channels=3, dtype=jnp.float32)
lr = jnp.ones((1, 128, 128, 3), dtype=jnp.float32)
params = model.init(PRNGKey(0), lr)
sr = model.apply(params, lr) #[1, 256, 256, 3]Not sure which models are currently the best? Check out the PapersWithCode Image Super-Resolution Leaderboards
- Perceptual Loss (VGG16)
- L1 Loss
- L2 Loss
- Peak-Signal-Noise-Ratio (PSNR) Loss
- Mean Squared Error (MSE)
- Mean Absolute Error (MAE)
- PSNR
- PixelShuffle
- Upsample
- Sequential
- PReLU