Update README.md

README.md

@@ -2,24 +2,31 @@
 tags:
 - colorization
 - image colorization
+- unet
+- resnet
+- coco
+- image-to-image
 license: apache-2.0
+thumbnail: https://cdn-uploads.huggingface.co/production/uploads/6318256d212fce5a3cde0fe3/To7FVLusBz1kl8g9HhPIf.png
 datasets:
 - ayushshah/coco-2017-image-colorization-224
 pipeline_tag: image-to-image
 ---
 
 # Image colorization
-Created a UNet architecture, utilizing transfer learning by using a pretrained ResNet-34 as an encoder.
+A UNet architecture, utilizing transfer learning by using a pretrained ResNet-34 as an encoder.
+
+Try the model on [Colab](https://colab.research.google.com/drive/1JYbSLtDuFSw2NYe-YW-kZHLNkt4-v7jd) or [Huggingface space](https://huggingface.co/spaces/ayushshah/imagecolorization).
 <img src="https://cdn-uploads.huggingface.co/production/uploads/6318256d212fce5a3cde0fe3/To7FVLusBz1kl8g9HhPIf.png" width="800px"/>
 
-The model takes a 224x224 L channel tensor as input and outputs 224x224 ab channels. There are no dedicated datasets for image colorisation, hence I curated my own dataset and used it to train the model. The dataset can be found [here](https://huggingface.co/datasets/ayushshah/coco-2017-image-colorization-224). This repository contains the model weights.
+The model takes a 1x224x224 L tensor as input and outputs 2x224x224 ab channels. The decoder has been trained from scratch. The encoder (ResNet-34) was initially frozen for the decoder to adapt to the task, then it was progressively unfreezed layer by layer. Initial layers were not unfreezed, only deeper layers were fine-tuned. Read various research papers. It took 20+ hours of training on Google Colab and Kaggle T4 GPUs to train the model. 
+
+There are no dedicated datasets for image colorisation, hence I curated my own dataset and used it to train the model. The COCO 2017 dataset was filtered to remove grayscale images, heavily filtered images, and other artifacts not suitable for training a natural colorization model. Also the images were center-cropped and resized to 224x224. The dataset can be found [here](https://huggingface.co/datasets/ayushshah/coco-2017-image-colorization-224). This repository contains the model weights and the UNet architecture to load the weights into.
 
-<div style="display: flex; justify-content: space-evenly;"> 
-  <img src="https://cdn-uploads.huggingface.co/production/uploads/6318256d212fce5a3cde0fe3/5eKOaiTUK4uDeq07MdJIY.png" width="650px"/>
-  <img src="https://cdn-uploads.huggingface.co/production/uploads/6318256d212fce5a3cde0fe3/MI08kcYgav2ouXlHvvNtu.png" width="650px"/>
-</div>
+<img src="https://cdn-uploads.huggingface.co/production/uploads/6318256d212fce5a3cde0fe3/5eKOaiTUK4uDeq07MdJIY.png" width="650px"/>
+<img src="https://cdn-uploads.huggingface.co/production/uploads/6318256d212fce5a3cde0fe3/MI08kcYgav2ouXlHvvNtu.png" width="650px"/>
 
-# Refereences
+# References
 - [Let there be Color!: Joint End-to-end Learning of Global and Local Image Priors for Automatic Image Colorization with Simultaneous Classification](https://iizuka.cs.tsukuba.ac.jp/projects/colorization/data/colorization_sig2016.pdf)
 - [Colorful Image Colorization](https://arxiv.org/pdf/1603.08511)
 - [Color and Attention for U: Modified Multi Attention U-Net for a Better Image Colorization](https://joiv.org/index.php/joiv/article/view/1828)