Home

Awesome

jax-enhance


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.

Installation

pip install git+https://github.com/isaaccorley/jax-enhance

Models

The following models are currently implemented:

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]

State-of-the-Art

Not sure which models are currently the best? Check out the PapersWithCode Image Super-Resolution Leaderboards

Losses

Metrics

Other layers not available in Flax