Add update to README with datasets training recipe.

README.md

@@ -5,6 +5,10 @@ tags:
 - image-super-resolution
 datasets:
 - eugenesiow/Div2k
+- eugenesiow/Set5
+- eugenesiow/Set14
+- eugenesiow/BSD100
+- eugenesiow/Urban100
 metrics:
 - pnsr
 - ssim
@@ -42,8 +46,9 @@ preds = model(inputs)
 ImageLoader.save_image(preds, './scaled_2x.png')                        # save the output 2x scaled image to `./scaled_2x.png`
 ImageLoader.save_compare(inputs, preds, './scaled_2x_compare.png')      # save an output comparing the super-image with a bicubic scaling
 ```
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/eugenesiow/super-image-notebooks/blob/master/notebooks/Upscale_Images_with_Pretrained_super_image_Models.ipynb "Open in Colab")
 ## Training data
-The models for 2x, 3x and 4x image super resolution were pretrained on [DIV2K](https://data.vision.ee.ethz.ch/cvl/DIV2K/), a dataset of 800 high-quality (2K resolution) images for training, augmented to 4000 images and uses a dev set of  100 validation images (images numbered 801 to 900). 
+The models for 2x, 3x and 4x image super resolution were pretrained on [DIV2K](https://huggingface.co/datasets/eugenesiow/Div2k), a dataset of 800 high-quality (2K resolution) images for training, augmented to 4000 images and uses a dev set of  100 validation images (images numbered 801 to 900). 
 ## Training procedure
 ### Preprocessing
 We follow the pre-processing and training method of [Wang et al.](https://arxiv.org/abs/2104.07566).
@@ -51,24 +56,20 @@ Low Resolution (LR) images are created by using bicubic interpolation as the res
 During training, RGB patches with size of 64×64 from the LR input are used together with their corresponding HR patches. 
 Data augmentation is applied to the training set in the pre-processing stage where five images are created from the four corners and center of the original image. 
 
-The following code provides some helper functions to preprocess the data.
+We need the huggingface [datasets](https://huggingface.co/datasets?filter=task_ids:other-other-image-super-resolution) library to download the data:
+```bash
+pip install datasets
+```
+The following code gets the data and preprocesses/augments the data.
+
 ```python
-from super_image.data import EvalDataset, TrainAugmentDataset, DatasetBuilder
+from datasets import load_dataset
+from super_image.data import EvalDataset, TrainDataset, augment_five_crop
 
-DatasetBuilder.prepare(
-    base_path='./DIV2K/DIV2K_train_HR',
-    output_path='./div2k_4x_train.h5',
-    scale=4,
-    do_augmentation=True
-)
-DatasetBuilder.prepare(
-    base_path='./DIV2K/DIV2K_val_HR',
-    output_path='./div2k_4x_val.h5',
-    scale=4,
-    do_augmentation=False
-)
-train_dataset = TrainAugmentDataset('./div2k_4x_train.h5', scale=4)
-val_dataset = EvalDataset('./div2k_4x_val.h5')
+augmented_dataset = load_dataset('eugenesiow/Div2k', 'bicubic_x4', split='train')\
+    .map(augment_five_crop, batched=True, desc="Augmenting Dataset")                                # download and augment the data with the five_crop method
+train_dataset = TrainDataset(augmented_dataset)                                                     # prepare the train dataset for loading PyTorch DataLoader
+eval_dataset = EvalDataset(load_dataset('eugenesiow/Div2k', 'bicubic_x4', split='validation'))      # prepare the eval dataset for the PyTorch DataLoader
 ```
 ### Pretraining
 The model was trained on GPU. The training code is provided below:
@@ -90,19 +91,21 @@ trainer = Trainer(
     model=model,                         # the instantiated model to be trained
     args=training_args,                  # training arguments, defined above
     train_dataset=train_dataset,         # training dataset
-    eval_dataset=val_dataset             # evaluation dataset
+    eval_dataset=eval_dataset            # evaluation dataset
 )
 
 trainer.train()
 ```
+
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/eugenesiow/super-image-notebooks/blob/master/notebooks/Train_super_image_Models.ipynb "Open in Colab")
 ## Evaluation results
 The evaluation metrics include [PSNR](https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio#Quality_estimation_with_PSNR) and [SSIM](https://en.wikipedia.org/wiki/Structural_similarity#Algorithm). 
 
 Evaluation datasets include:
-- Set5 - [Bevilacqua et al. (2012)](http://people.rennes.inria.fr/Aline.Roumy/results/SR_BMVC12.html)
-- Set14 - [Zeyde et al. (2010)](https://sites.google.com/site/romanzeyde/research-interests)
-- BSD100 - [Martin et al. (2001)](https://www.eecs.berkeley.edu/Research/Projects/CS/vision/bsds/)
-- Urban100 - [Huang et al. (2015)](https://sites.google.com/site/jbhuang0604/publications/struct_sr)
+- Set5 - [Bevilacqua et al. (2012)](https://huggingface.co/datasets/eugenesiow/Set5)
+- Set14 - [Zeyde et al. (2010)](https://huggingface.co/datasets/eugenesiow/Set14)
+- BSD100 - [Martin et al. (2001)](https://huggingface.co/datasets/eugenesiow/BSD100)
+- Urban100 - [Huang et al. (2015)](https://huggingface.co/datasets/eugenesiow/Urban100)
 
 The results columns below are represented below as `PSNR/SSIM`. They are compared against a Bicubic baseline.
 
@@ -122,6 +125,9 @@ The results columns below are represented below as `PSNR/SSIM`. They are compare
 |Urban100  	    |4x  	    |23.14/0.6573  	    |**26.10/0.7857**  	    |
 
 ![Comparing Bicubic upscaling against the models x4 upscaling on Set5 Image 2](images/msrn_2_4_compare.png "Comparing Bicubic upscaling against the models x4 upscaling on Set5 Image 2")
+You can find a notebook to easily run evaluation on pretrained models below:
+
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/eugenesiow/super-image-notebooks/blob/master/notebooks/Evaluate_Pretrained_super_image_Models.ipynb "Open in Colab")
 
 ## BibTeX entry and citation info
 ```bibtex