Skip to content
/ zipnet Public

Neural compression algorithm run entirely on the web.

License

Notifications You must be signed in to change notification settings

Conzel/zipnet

Repository files navigation

zipnet

Zipnet is the first useable neural image encoder and decoder that can run entirely in your web browser (no servers involved)!

Useable means:

  • Superior performance compared to classical codecs (JPEG)
  • Acceptable encoding and decoding speed (a few seconds)

This is achieved by defining a neural codec in Pytorch as well as in Rust, running training in Python and transferring the trained weights to Rust for platform-independent inference. The codec is compiled to WebASM, which has inference runtimes roughly compared to native code, and is then included in our website.

We have re-written a lot neural network modules in pure Rust for maximum portability. These can be found under src/ml, or in our crate blowtorch, which serves as a framework for porting Pytorch weights to Rust for inference.

This project has spun out of a course project for the lecture Data Compression With Deep Probabilistic Models offered by Robert Bamler at the University of Tübingen.

If you are interested in what all of the crates and modules do, the folders either have READMEs or you can refer to the compiled cargo doc. You can compile the documentation for all crates you are interested in by switching to the corresponding folder (e.g. cd coders) and running cargo doc --open.

Using the CLI

  • Install the Rust toolchain (https://www.rust-lang.org/tools/install)
  • Run the CLI via cargo run -p zipnet -- <all the other commands you want to add>. You can get an overview over the possibilites by running cargo run -p zipnet -- --help. For subcommands, just call the subcommand, and then --help to get information on them. Alternatively, you can build the program via cargo build --release (debug is likely too slow) and access the binary under target/release/zipnet.

In the following, we present some example usages.

Compression

./zipnet compress /path/to/image For compression, we accept .jpg, .png and .npy as inputs (.npy inputs are not preprocessed and are mainly intended for debugging). The compressed file is written out to /path/to/image.bin, overwriting the output file if it is already present.

Decompression

./zipnet decompress /path/to/image.bin /path/to/output_image.png Decompression takes a previous output from the compression step and outputs a recovered version of the image under the given path. The output is always a .png.

Debugging with autoencoder

./zipnet autoencoder /path/to/image.png The output is under /path/to/image-reconstructed-ae.png. The autoencoder does not perform any quantization, thus the image quality will be a lot better. This is used mainly for debugging. Note that the final image will have unspecified dimensions and has a small black border in most cases – this is not present in the image produced with the encoder-decoder pair.

Building WASM locally

Setup (Ubuntu)

Make sure you have installed the Rust-toolchain, wasm-pack and npm.

  • Rust: Follow the instructions at https://www.rust-lang.org/tools/install
  • wasm-pack: Run curl https://rustwasm.github.io/wasm-pack/installer/init.sh -sSf | sh
  • npm: If you don't have npm installed, run sudo apt install nodejs npm. If you have it installed, make you have the latest version by running npm install npm@latest -g

Build

You can build each target specified in the Cargo.toml under workspace-members by running cargo build -p <target-name>.

To setup the website locally, follow these steps from the root directory:

  • cd wasm
  • wasm-pack build
  • cd www
  • npm install
  • npm run start
  • Open your browser and navigate to http://localhost:8080/

You should see the website running! As long as you have the server running, you can make changes by just re-running wasm-pack build.

You can get more detailed info on why we are doing everything like that at https://rustwasm.github.io/book/game-of-life/hello-world.html.

Testing

As always, tests can be run via cargo test. There exist automated, randomized tests for the Rust modules that implement Python functionality to closely match their behaviour. They can be generated through the script by simply running the script in scripts/generate_tests.py.

For this, we recommend tensorflow 2.5.0, numpy 1.19.5 and Jinja2 3.0.0.

List of crates tested for WASM-compatibility

  • ndarray v0.15.3 (no feature flags)

Contributing for maintainers

Setting up commit hooks

Please ensure you have git hooks installed when you push changes. We use rusty-hook to manage git-hooks. To install, run cargo install rusty-hook and then rusty-hook init in this directory. We run cargo fmt and cargo check on commits.

General note

If you make changes and push them to the main branch, they are automatically reflected on the website via a Github Action. Thus please ensure beforehand that everything works locally as you want it. Please also make sure, that the GitHub Action correctly runs and fix errors that can occur.