unet.cpp

U-Net from scratch in C++. Based on milesial/Pytorch-UNet.

|

Instructions

1. Download and convert weights

python convert.py

Downloaded weights come from here. Weights are converted to llama.c format as weights.bin.

2. Build U-Net

make unet

3. Run U-Net

./unet [input] [output]

Input must be an image in PPM format with spatial dimensions powers of 2. Note that the current implementation of the PPM reader is limited; however, it has worked successfully on files created with ImageMagick.

The output will be saved in PGM format.

Example

A sample valid PPM file is included as car.ppm (created with ImageMagick from car.png via convert car.png car.ppm).

./unet car.ppm mask.pgm

Future work

unet.cpp is slow! The matrix multiply underlying the convolutions (via im2col) could stand to be optimized, but more likely a better convolution algorithm will be needed (FFT or Winograd).