[WIP] U-Net

[alexander-g]

• U-Net architecture for image segmentation

• MeanIoU metric

• class_weight and ignore_index parameters for SparseCategoricalCrossentropy

• Backboned U-Net

• Training examples

• Pretrained model parameters

• mkdocs

[cgarciae]

Looking good! Hey I am working on making Elegy layers numerically equivalent with Flax layers in #128, I am just starting with Linear tho, but eventually this could lead to us being able to reuse their pre-trained models and vice-versa.

[alexander-g]

• This PR is finished but depends on PR Module Slicing #115 so I'm not marking it as ready yet
• Includes UNets with ResNet18 and ResNet50 backbones pretrained on COCO with 384px resolution. Not entirely happy with it yet but if I go up the resolution I need to reduce the batch size and the performance drops. I will try again some time later.
• @cgarciae Please take a look at module.py:171 whether or not this breaks anything

fname,_  = urllib.request.urlretrieve('https://upload.wikimedia.org/wikipedia/commons/thumb/3/3d/Take_ours%21.jpg/800px-Take_ours%21.jpg')
x        = np.array(PIL.Image.open(fname).resize([384,384])) / np.float32(255)
unet_r18 = elegy.nets.unet.UNet_R18(weights='coco')
y        = elegy.Model(unet_r18).predict(x[np.newaxis])[0]
i        = elegy.nets.unet.COCO_CLASSNAMES_TO_CHANNELS['sheep']

figure(0, (10,4))
subplot(121); imshow(x);               axis('off');
subplot(122); imshow(y.argmax(-1)==i); axis('off');