Eugene Siow commited on
Commit
5d7409d
1 Parent(s): d6b8d33

Add new training recipe.

Browse files
Files changed (1) hide show
  1. README.md +33 -26
README.md CHANGED
@@ -5,6 +5,10 @@ tags:
5
  - image-super-resolution
6
  datasets:
7
  - eugenesiow/Div2k
 
 
 
 
8
  metrics:
9
  - pnsr
10
  - ssim
@@ -33,15 +37,16 @@ import requests
33
  url = 'https://paperswithcode.com/media/datasets/Set5-0000002728-07a9793f_zA3bDjj.jpg'
34
  image = Image.open(requests.get(url, stream=True).raw)
35
 
36
- model = MsrnModel.from_pretrained('eugenesiow/msrn', scale=2) # scale 2, 3 and 4 models available
37
  inputs = ImageLoader.load_image(image)
38
  preds = model(inputs)
39
 
40
- ImageLoader.save_image(preds, './scaled_2x.png') # save the output 2x scaled image to `./scaled_2x.png`
41
- ImageLoader.save_compare(inputs, preds, './scaled_2x_compare.png') # save an output comparing the super-image with a bicubic scaling
42
  ```
 
43
  ## Training data
44
- 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).
45
  ## Training procedure
46
  ### Preprocessing
47
  We follow the pre-processing and training method of [Wang et al.](https://arxiv.org/abs/2104.07566).
@@ -49,24 +54,20 @@ Low Resolution (LR) images are created by using bicubic interpolation as the res
49
  During training, RGB patches with size of 64×64 from the LR input are used together with their corresponding HR patches.
50
  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.
51
 
52
- The following code provides some helper functions to preprocess the data.
 
 
 
 
 
53
  ```python
54
- from super_image.data import EvalDataset, TrainAugmentDataset, DatasetBuilder
 
55
 
56
- DatasetBuilder.prepare(
57
- base_path='./DIV2K/DIV2K_train_HR',
58
- output_path='./div2k_4x_train.h5',
59
- scale=4,
60
- do_augmentation=True
61
- )
62
- DatasetBuilder.prepare(
63
- base_path='./DIV2K/DIV2K_val_HR',
64
- output_path='./div2k_4x_val.h5',
65
- scale=4,
66
- do_augmentation=False
67
- )
68
- train_dataset = TrainAugmentDataset('./div2k_4x_train.h5', scale=4)
69
- val_dataset = EvalDataset('./div2k_4x_val.h5')
70
  ```
71
  ### Pretraining
72
  The model was trained on GPU. The training code is provided below:
@@ -87,23 +88,25 @@ trainer = Trainer(
87
  model=model, # the instantiated model to be trained
88
  args=training_args, # training arguments, defined above
89
  train_dataset=train_dataset, # training dataset
90
- eval_dataset=val_dataset # evaluation dataset
91
  )
92
 
93
  trainer.train()
94
  ```
 
 
95
  ## Evaluation results
96
  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).
97
 
98
  Evaluation datasets include:
99
- - Set5 - [Bevilacqua et al. (2012)](http://people.rennes.inria.fr/Aline.Roumy/results/SR_BMVC12.html)
100
- - Set14 - [Zeyde et al. (2010)](https://sites.google.com/site/romanzeyde/research-interests)
101
- - BSD100 - [Martin et al. (2001)](https://www.eecs.berkeley.edu/Research/Projects/CS/vision/bsds/)
102
- - Urban100 - [Huang et al. (2015)](https://sites.google.com/site/jbhuang0604/publications/struct_sr)
103
 
104
  The results columns below are represented below as `PSNR/SSIM`. They are compared against a Bicubic baseline.
105
 
106
- |Dataset |Scale |Bicubic |msrn-bam |
107
  |--- |--- |--- |--- |
108
  |Set5 |2x |33.64/0.9292 | |
109
  |Set5 |3x |30.39/0.8678 | |
@@ -120,6 +123,10 @@ The results columns below are represented below as `PSNR/SSIM`. They are compare
120
 
121
  ![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")
122
 
 
 
 
 
123
  ## BibTeX entry and citation info
124
  ```bibtex
125
  @InProceedings{Agustsson_2017_CVPR_Workshops,
 
5
  - image-super-resolution
6
  datasets:
7
  - eugenesiow/Div2k
8
+ - eugenesiow/Set5
9
+ - eugenesiow/Set14
10
+ - eugenesiow/BSD100
11
+ - eugenesiow/Urban100
12
  metrics:
13
  - pnsr
14
  - ssim
 
37
  url = 'https://paperswithcode.com/media/datasets/Set5-0000002728-07a9793f_zA3bDjj.jpg'
38
  image = Image.open(requests.get(url, stream=True).raw)
39
 
40
+ model = MsrnModel.from_pretrained('eugenesiow/msrn', scale=4) # scale 2, 3 and 4 models available
41
  inputs = ImageLoader.load_image(image)
42
  preds = model(inputs)
43
 
44
+ ImageLoader.save_image(preds, './scaled_4x.png') # save the output 4x scaled image to `./scaled_4x.png`
45
+ ImageLoader.save_compare(inputs, preds, './scaled_4x_compare.png') # save an output comparing the super-image with a bicubic scaling
46
  ```
47
+ [![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")
48
  ## Training data
49
+ 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).
50
  ## Training procedure
51
  ### Preprocessing
52
  We follow the pre-processing and training method of [Wang et al.](https://arxiv.org/abs/2104.07566).
 
54
  During training, RGB patches with size of 64×64 from the LR input are used together with their corresponding HR patches.
55
  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.
56
 
57
+ We need the huggingface [datasets](https://huggingface.co/datasets?filter=task_ids:other-other-image-super-resolution) library to download the data:
58
+ ```bash
59
+ pip install datasets
60
+ ```
61
+ The following code gets the data and preprocesses/augments the data.
62
+
63
  ```python
64
+ from datasets import load_dataset
65
+ from super_image.data import EvalDataset, TrainDataset, augment_five_crop
66
 
67
+ augmented_dataset = load_dataset('eugenesiow/Div2k', 'bicubic_x4', split='train')\
68
+ .map(augment_five_crop, batched=True, desc="Augmenting Dataset") # download and augment the data with the five_crop method
69
+ train_dataset = TrainDataset(augmented_dataset) # prepare the train dataset for loading PyTorch DataLoader
70
+ eval_dataset = EvalDataset(load_dataset('eugenesiow/Div2k', 'bicubic_x4', split='validation')) # prepare the eval dataset for the PyTorch DataLoader
 
 
 
 
 
 
 
 
 
 
71
  ```
72
  ### Pretraining
73
  The model was trained on GPU. The training code is provided below:
 
88
  model=model, # the instantiated model to be trained
89
  args=training_args, # training arguments, defined above
90
  train_dataset=train_dataset, # training dataset
91
+ eval_dataset=eval_dataset # evaluation dataset
92
  )
93
 
94
  trainer.train()
95
  ```
96
+
97
+ [![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")
98
  ## Evaluation results
99
  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).
100
 
101
  Evaluation datasets include:
102
+ - Set5 - [Bevilacqua et al. (2012)](https://huggingface.co/datasets/eugenesiow/Set5)
103
+ - Set14 - [Zeyde et al. (2010)](https://huggingface.co/datasets/eugenesiow/Set14)
104
+ - BSD100 - [Martin et al. (2001)](https://huggingface.co/datasets/eugenesiow/BSD100)
105
+ - Urban100 - [Huang et al. (2015)](https://huggingface.co/datasets/eugenesiow/Urban100)
106
 
107
  The results columns below are represented below as `PSNR/SSIM`. They are compared against a Bicubic baseline.
108
 
109
+ |Dataset |Scale |Bicubic |msrn |
110
  |--- |--- |--- |--- |
111
  |Set5 |2x |33.64/0.9292 | |
112
  |Set5 |3x |30.39/0.8678 | |
 
123
 
124
  ![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")
125
 
126
+ You can find a notebook to easily run evaluation on pretrained models below:
127
+
128
+ [![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")
129
+
130
  ## BibTeX entry and citation info
131
  ```bibtex
132
  @InProceedings{Agustsson_2017_CVPR_Workshops,