Add training and inference code samples.

README.md

@@ -27,14 +27,19 @@ pip install super-image
 ```
 Here is how to use a pre-trained model to upscale your image:
 ```python
-from super_image import EdsrModel, DataLoader
+from super_image import EdsrModel, ImageLoader
 from PIL import Image
 import requests
-url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
+
+url = 'https://paperswithcode.com/media/datasets/Set5-0000002728-07a9793f_zA3bDjj.jpg'
 image = Image.open(requests.get(url, stream=True).raw)
-model = EdsrModel.from_pretrained('edsr')
-inputs = DataLoader(images=image)
+
+model = EdsrModel.from_pretrained('eugenesiow/edsr-base', scale=2)
+inputs = ImageLoader.load_image(image)
 preds = model(inputs)
+
+ImageLoader.save_image(preds, './scaled_2x.png')
+ImageLoader.save_compare(inputs, preds, './scaled_2x_compare.png')
 ```
 ## 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). 
@@ -44,17 +49,40 @@ We follow the pre-processing and training method of [Wang et. al](https://arxiv.
 Low Resolution (LR) images are created by using bicubic interpolation as the resizing method to reduce the size of the High Resolution (HR) images by x2, x3 and x4 times.
 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.
+```python
+from super_image.data import EvalDataset, TrainAugmentDataset, DatasetBuilder
+
+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')
+```
 ### Pretraining
 The model was trained on GPU. The training code is provided below:
 ```python
 from super_image import Trainer, TrainingArguments, EdsrModel, EdsrConfig
 
 training_args = TrainingArguments(
-    output_dir='./results',               # output directory
-    num_train_epochs=1000,                # total number of training epochs
+    output_dir='./results',                 # output directory
+    num_train_epochs=1000,                  # total number of training epochs
 )
 
-config = EdsrConfig()
+config = EdsrConfig(
+    scale=4,                                # train a model to upscale 4x
+    supported_scales=[2, 3, 4],
+)
 model = EdsrModel(config)
 
 trainer = Trainer(