Add hf 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 EDSR 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 EDSR 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,14 @@ 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.
+The following code provides some helper functions to get the data and preprocess/augment the data.
 ```python
-from super_image.data import EvalDataset, TrainAugmentDataset, DatasetBuilder
+from datasets import load_dataset
 
-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:
@@ -120,7 +115,7 @@ The results columns below are represented below as `PSNR/SSIM`. They are compare
 |Urban100  	    |3x  	    |  	                |**29.23/0.8723**  	    |
 |Urban100  	    |4x  	    |23.14/0.6573  	    |**26.02/0.7832**  	    |
 
-![Comparing Bicubic upscaling against EDSR x2 upscaling on Set5 Image 2](images/Set5_2_compare.png "Comparing Bicubic upscaling against EDSR x2 upscaling on Set5 Image 2")
+![Comparing Bicubic upscaling against x2 upscaling on Set5 Image 2](images/Set5_2_compare.png "Comparing Bicubic upscaling against x2 upscaling on Set5 Image 2")
 
 ## BibTeX entry and citation info
 ```bibtex