init model

README.md

@@ -0,0 +1,121 @@
+---
+license: apache-2.0
+datasets:
+- eugenesiow/Div2k
+- eugenesiow/Set5
+language:
+- en
+tags:
+- RyzenAI
+- super resolution
+- SISR
+- pytorch
+---
+## Model description
+SESR is based on linear overparameterization of CNNs and creates an efficient model architecture for SISR. It was introduced in the paper [Collapsible Linear Blocks for Super-Efficient Super Resolution](https://arxiv.org/abs/2103.09404).
+The official code for this work is available at this 
+https://github.com/ARM-software/sesr
+
+We develop a modified version that could be supported by [AMD Ryzen AI](https://onnxruntime.ai/docs/execution-providers/Vitis-AI-ExecutionProvider.html).
+
+## Intended uses & limitations
+
+You can use the raw model for super resolution. See the [model hub](https://huggingface.co/models?search=amd/sesr) to look for all available models.
+
+
+## How to use
+
+### Installation
+
+   Follow [Ryzen AI Installation](https://ryzenai.docs.amd.com/en/latest/inst.html) to prepare the environment for Ryzen AI.
+   Run the following script to install pre-requisites for this model.
+   ```bash
+   pip install -r requirements.txt 
+   ```
+
+
+### Data Preparation (optional: for accuracy evaluation)
+
+1. Download the benchmark(https://cv.snu.ac.kr/research/EDSR/benchmark.tar) dataset.
+2. Organize the dataset directory as follows:
+```Plain
+└── dataset
+     └── benchmark
+          ├── Set5
+               ├── HR
+               |   ├── baby.png
+               |   ├── ...
+               └── LR_bicubic
+                   └──X2
+                      ├──babyx2.png
+                      ├── ...
+          ├── Set14
+          ├── ...    
+```
+
+### Test & Evaluation
+
+- Code snippet from [`one_image_inference.py`](one_image_inference.py) on how to use
+```python
+    parser = argparse.ArgumentParser(description='EDSR and MDSR')
+    parser.add_argument('--onnx_path', type=str, default='SESR_int8.onnx',
+                    help='onnx path')
+    parser.add_argument('--image_path', default='test_data/test.png',
+                    help='path of your image')
+    parser.add_argument('--output_path', default='test_data/sr.png',
+                    help='path of your image')
+    parser.add_argument('--ipu', action='store_true',
+                    help='use ipu')
+    parser.add_argument('--provider_config', type=str, default=None,
+                    help='provider config path')
+    args = parser.parse_args()
+    if args.ipu:
+        providers = ["VitisAIExecutionProvider"]
+        provider_options = [{"config_file": args.provider_config}]
+    else:
+        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+        provider_options = None
+   
+    onnx_file_name = args.onnx_path
+    image_path = args.image_path
+    output_path = args.output_path
+
+    ort_session = onnxruntime.InferenceSession(onnx_file_name,  providers=providers, provider_options=provider_options) 
+    lr = cv2.imread(image_path)[np.newaxis,:,:,:].transpose((0,3,1,2)).astype(np.float32)
+    sr = tiling_inference(ort_session, lr, 8, (56, 56))
+    sr = np.clip(sr, 0, 255)
+    sr = sr.squeeze().transpose((1,2,0)).astype(np.uint8)
+    sr = cv2.imwrite(output_path, sr)
+```
+
+ - Run inference for a single image
+  ```python
+  python one_image_inference.py --onnx_path SESR_int8.onnx --image_path /Path/To/Your/Image --ipu --provider_config Path/To/vaip_config.json
+  ```
+Note: **vaip_config.json**  is located at the setup package of Ryzen AI (refer to [Installation](https://huggingface.co/amd/yolox-s#installation))
+
+ - Test accuracy of the quantized model
+  ```python
+  python test.py --onnx_path SESR_int8.onnx --data_test Set5 --ipu --provider_config Path/To/vaip_config.json 
+  ```
+
+
+
+### Performance
+| Method     | Scale | Flops | Set5         |
+|------------|-------|-------|--------------|
+|SESR-S (float) |X2 |10.22G |37.21|
+|SESR-S (INT8) |X2  |10.22G |36.81|
+- Note: the Flops is calculated with the input resolution is 256x256
+
+
+```bibtex
+@misc{bhardwaj2022collapsible,
+      title={Collapsible Linear Blocks for Super-Efficient Super Resolution}, 
+      author={Kartikeya Bhardwaj and Milos Milosavljevic and Liam O'Neil and Dibakar Gope and Ramon Matas and Alex Chalfin and Naveen Suda and Lingchuan Meng and Danny Loh},
+      year={2022},
+      eprint={2103.09404},
+      archivePrefix={arXiv},
+      primaryClass={eess.IV}
+}
+```

SESR_int8.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2b4dc9547653f01bec1ba53e42f6722ffa0cc74ee1c787c7e93174d17260de7
+size 110994

data/__init__.py

@@ -0,0 +1,36 @@
+from importlib import import_module
+#from dataloader import MSDataLoader
+from torch.utils.data import dataloader
+from torch.utils.data import ConcatDataset
+import torch
+import random
+# This is a simple wrapper function for ConcatDataset
+class MyConcatDataset(ConcatDataset):
+    def __init__(self, datasets):
+        super(MyConcatDataset, self).__init__(datasets)
+      
+
+    def set_scale(self, idx_scale):
+        for d in self.datasets:
+            if hasattr(d, 'set_scale'): d.set_scale(idx_scale)
+
+class Data:
+    def __init__(self, args):
+        self.loader_train = None
+        self.loader_test = []
+        for d in args.data_test:
+            if d in ['Set5', 'Set14', 'B100', 'Urban100']:
+                m = import_module('data.benchmark')
+                testset = getattr(m, 'Benchmark')(args, name=d)
+            else:
+                assert NotImplementedError 
+
+            self.loader_test.append(
+                dataloader.DataLoader(
+                    testset,
+                    batch_size=1,
+                    shuffle=False,
+                    pin_memory=False,
+                    num_workers=args.n_threads,
+                )
+            )

data/benchmark.py

@@ -0,0 +1,23 @@
+import os
+
+#from data import common
+from data import srdata
+import numpy as np
+import torch
+import torch.utils.data as data
+
+class Benchmark(srdata.SRData):
+    def __init__(self, args, name='',  benchmark=True):
+        super(Benchmark, self).__init__(
+            args, name=name,  benchmark=True
+        )
+
+    def _set_filesystem(self, dir_data):
+        self.apath = os.path.join(dir_data, 'benchmark', self.name)
+        self.dir_hr = os.path.join(self.apath, 'HR')
+        if self.input_large:
+            self.dir_lr = os.path.join(self.apath, 'LR_bicubicL')
+        else:
+            self.dir_lr = os.path.join(self.apath, 'LR_bicubic')
+        self.ext = ('', '.png')
+

data/common.py

@@ -0,0 +1,34 @@
+import random
+
+import numpy as np
+import skimage.color as sc
+
+import torch
+
+def set_channel(*args, n_channels=3):
+    def _set_channel(img):
+        if img.ndim == 2:
+            img = np.expand_dims(img, axis=2)
+
+        c = img.shape[2]
+        if n_channels == 1 and c == 3:
+            img = np.expand_dims(sc.rgb2ycbcr(img)[:, :, 0], 2)
+        elif n_channels == 3 and c == 1:
+            img = np.concatenate([img] * n_channels, 2)
+
+        return img
+
+    return [_set_channel(a) for a in args]
+
+def np2Tensor(*args, rgb_range=255):
+    def _np2Tensor(img):
+        np_transpose = np.ascontiguousarray(img.transpose((2, 0, 1)))
+        tensor = torch.from_numpy(np_transpose).float()
+        tensor.mul_(rgb_range / 255)
+
+        return tensor
+
+    return [_np2Tensor(a) for a in args]
+
+
+

data/data_tiling.py

@@ -0,0 +1,46 @@
+import numpy as np
+import math
+
+
+def tiling_inference(session, lr, overlapping, patch_size):
+    """
+    Parameters: 
+        - session: an ONNX Runtime session object that contains the super-resolution model 
+        - lr: the low-resolution image 
+        - overlapping: the number of pixels to overlap between adjacent patches 
+        - patch_size: a tuple of (height, width) that specifies the size of each patch
+    Returns: - a numpy array that represents the enhanced image
+    """
+    _, _, h, w = lr.shape
+    sr = np.zeros((1, 3, 2*h, 2*w))
+    n_h = math.ceil(h / float(patch_size[0] - overlapping))
+    n_w = math.ceil(w / float(patch_size[1] - overlapping))
+    #every tilling input has same size of patch_size
+    for ih in range(n_h):
+        h_idx = ih * (patch_size[0] - overlapping)
+        h_idx = h_idx if h_idx + patch_size[0] <= h else h - patch_size[0]
+        for iw in range(n_w):
+            w_idx = iw * (patch_size[1] - overlapping)
+            w_idx = w_idx if w_idx + patch_size[1] <= w else w - patch_size[1]
+
+            tilling_lr = lr[..., h_idx: h_idx+patch_size[0], w_idx: w_idx+patch_size[1]]
+            sr_tiling = session.run(None, {session.get_inputs()[0].name: tilling_lr})[0]
+
+            left, right, top, bottom = 0, patch_size[1], 0, patch_size[0]
+            left += overlapping//2 
+            right -= overlapping//2
+            top += overlapping//2
+            bottom -= overlapping//2
+            #processing edge pixels
+            if w_idx == 0:
+                left -= overlapping//2
+            if h_idx == 0:
+                top -= overlapping//2
+            if h_idx+patch_size[0]>=h:
+                bottom += overlapping//2
+            if w_idx+patch_size[1]>=w:
+                right += overlapping//2
+            
+            #get preditions
+            sr[... , 2*(h_idx+top): 2*(h_idx+bottom), 2*(w_idx+left): 2*(w_idx+right)] = sr_tiling[..., 2*top:2*bottom, 2*left:2*right]   
+    return sr

data/srdata.py

@@ -0,0 +1,81 @@
+import os
+import glob
+import random
+import pickle
+from data import common
+import imageio
+import torch.utils.data as data
+
+class SRData(data.Dataset):
+    def __init__(self, args, name='', benchmark=False):
+        self.args = args
+        self.name = name
+        self.split = 'test'
+        self.do_eval = True
+        self.benchmark = benchmark
+        self.input_large = False
+        self.scale = args.scale
+        self.idx_scale = 0
+        self._set_filesystem(args.dir_data)
+        list_hr, list_lr = self._scan()
+        self.images_hr, self.images_lr = list_hr, list_lr
+
+    # Below functions as used to prepare images
+    def _scan(self):
+        names_hr = sorted(
+            glob.glob(os.path.join(self.dir_hr, '*' + self.ext[0]))
+        )
+        names_lr = [[] for _ in self.scale]
+        for f in names_hr:
+            filename, _ = os.path.splitext(os.path.basename(f))
+            for si, s in enumerate(self.scale):
+                names_lr[si].append(os.path.join(
+                    self.dir_lr, 'X{}/{}x{}{}'.format(
+                        s, filename, s, self.ext[1]
+                    )
+                ))
+
+        return names_hr, names_lr
+
+    def _set_filesystem(self, dir_data):
+        self.apath = os.path.join(dir_data, self.name)
+        self.dir_hr = os.path.join(self.apath, 'HR')
+        self.dir_lr = os.path.join(self.apath, 'LR_bicubic')
+        self.ext = ('.png', '.png')
+
+    def __getitem__(self, idx):
+        lr, hr, filename = self._load_file(idx)
+        pair = self.get_patch(lr, hr)
+        pair = common.set_channel(*pair, n_channels=3)
+        pair_t = common.np2Tensor(*pair, rgb_range=255)
+
+        return pair_t[0], pair_t[1], filename
+
+    def __len__(self):
+        return len(self.images_hr)
+
+    def _get_index(self, idx):
+        return idx
+
+    def _load_file(self, idx):
+        idx = self._get_index(idx)
+        f_hr = self.images_hr[idx]
+        f_lr = self.images_lr[self.idx_scale][idx]
+
+        filename, _ = os.path.splitext(os.path.basename(f_hr))     
+        hr = imageio.imread(f_hr)
+        lr = imageio.imread(f_lr)
+        return lr, hr, filename
+
+    def get_patch(self, lr, hr):
+        scale = self.scale[self.idx_scale]
+        ih, iw = lr.shape[:2]
+        hr = hr[0:ih * scale, 0:iw * scale]
+        return lr, hr
+
+    def set_scale(self, idx_scale):
+        if not self.input_large:
+            self.idx_scale = idx_scale
+        else:
+            self.idx_scale = random.randint(0, len(self.scale) - 1)
+

metric.py

@@ -0,0 +1,83 @@
+import math
+import numpy as np
+from scipy import signal
+
+def calc_psnr(sr, hr, scale, rgb_range, benchmark=False):
+    if sr.size(-2) > hr.size(-2) or sr.size(-1) > hr.size(-1):
+        print("the dimention of sr image is not equal to hr's! ")
+        sr = sr[:,:,:hr.size(-2),:hr.size(-1)]
+    diff = (sr - hr).data.div(rgb_range)
+
+    if benchmark:
+        shave = scale
+        if diff.size(1) > 1:
+            convert = diff.new(1, 3, 1, 1)
+            convert[0, 0, 0, 0] = 65.738
+            convert[0, 1, 0, 0] = 129.057
+            convert[0, 2, 0, 0] = 25.064
+            diff.mul_(convert).div_(256)
+            diff = diff.sum(dim=1, keepdim=True)
+    else:
+        shave = scale + 6
+    valid = diff[:, :, shave:-shave, shave:-shave]
+    mse = valid.pow(2).mean()
+
+    return -10 * math.log10(mse)
+
+def matlab_style_gauss2D(shape=(3,3),sigma=0.5):
+    """
+    2D gaussian mask - should give the same result as MATLAB's fspecial('gaussian',[shape],[sigma])
+    Acknowledgement : https://stackoverflow.com/questions/17190649/how-to-obtain-a-gaussian-filter-in-python (Author@ali_m)
+    """
+    m,n = [(ss-1.)/2. for ss in shape]
+    y,x = np.ogrid[-m:m+1,-n:n+1]
+    h = np.exp( -(x*x + y*y) / (2.*sigma*sigma) )
+    h[ h < np.finfo(h.dtype).eps*h.max() ] = 0
+    sumh = h.sum()
+    if sumh != 0:
+      h /= sumh
+    return h
+
+def calc_ssim(X, Y, scale, rgb_range, dataset=None, sigma=1.5, K1=0.01, K2=0.03, R=255):
+    '''
+    X : y channel (i.e., luminance) of transformed YCbCr space of X
+    Y : y channel (i.e., luminance) of transformed YCbCr space of Y
+    Please follow the setting of psnr_ssim.m in EDSR (Enhanced Deep Residual Networks for Single Image Super-Resolution CVPRW2017).
+    Official Link : https://github.com/LimBee/NTIRE2017/tree/db34606c2844e89317aac8728a2de562ef1f8aba
+    The authors of EDSR use MATLAB's ssim as the evaluation tool, 
+    thus this function is the same as ssim.m in MATLAB with C(3) == C(2)/2. 
+    '''
+    gaussian_filter = matlab_style_gauss2D((11, 11), sigma) 
+    shave = scale
+    if X.size(1) > 1:
+        gray_coeffs = [65.738, 129.057, 25.064]
+        convert = X.new_tensor(gray_coeffs).view(1, 3, 1, 1) / 256
+        X = X.mul(convert).sum(dim=1)
+        Y = Y.mul(convert).sum(dim=1)
+
+
+    X = X[..., shave:-shave, shave:-shave].squeeze().cpu().numpy().astype(np.float64) 
+    Y = Y[..., shave:-shave, shave:-shave].squeeze().cpu().numpy().astype(np.float64)
+
+    window = gaussian_filter
+
+    ux = signal.convolve2d(X, window, mode='same', boundary='symm')
+    uy = signal.convolve2d(Y, window, mode='same', boundary='symm')
+
+    uxx = signal.convolve2d(X*X, window, mode='same', boundary='symm')
+    uyy = signal.convolve2d(Y*Y, window, mode='same', boundary='symm')
+    uxy = signal.convolve2d(X*Y, window, mode='same', boundary='symm')
+
+    vx = uxx - ux * ux
+    vy = uyy - uy * uy
+    vxy = uxy - ux * uy
+
+    C1 = (K1 * R) ** 2
+    C2 = (K2 * R) ** 2
+
+    A1, A2, B1, B2 = ((2 * ux * uy + C1, 2 * vxy + C2, ux ** 2 + uy ** 2 + C1, vx + vy + C2))
+    D = B1 * B2
+    S = (A1 * A2) / D
+    mssim = S.mean()
+
+    return mssim

one_image_inference.py

@@ -0,0 +1,44 @@
+import onnxruntime
+import cv2
+import numpy as np
+import sys
+import pathlib
+CURRENT_DIR = pathlib.Path(__file__).parent
+sys.path.append(str(CURRENT_DIR))
+from data.data_tiling import tiling_inference
+import argparse
+
+
+def main(args):
+    if args.ipu:
+        providers = ["VitisAIExecutionProvider"]
+        provider_options = [{"config_file": args.provider_config}]
+    else:
+        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+        provider_options = None
+    onnx_file_name = args.onnx_path
+    image_path = args.image_path
+    output_path = args.output_path
+
+    ort_session = onnxruntime.InferenceSession(onnx_file_name,  providers=providers, provider_options=provider_options) 
+    lr = cv2.imread(image_path)[np.newaxis,:,:,:].transpose((0,3,1,2)).astype(np.float32)
+    sr = tiling_inference(ort_session, lr, 8, (56, 56))
+    sr = np.clip(sr, 0, 255)
+    sr = sr.squeeze().transpose((1,2,0)).astype(np.uint8)
+    cv2.imwrite(output_path, sr)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='EDSR and MDSR')
+    parser.add_argument('--onnx_path', type=str, default='SESR_int8.onnx',
+                    help='onnx path')
+    parser.add_argument('--image_path', default='test_data/test.png',
+                    help='path of your image')
+    parser.add_argument('--output_path', default='test_data/sr.png',
+                    help='path of your image')
+    parser.add_argument('--ipu', action='store_true',
+                    help='use ipu')
+    parser.add_argument('--provider_config', type=str, default=None,
+                    help='provider config path')
+    args = parser.parse_args()
+    main(args)

option.py

@@ -0,0 +1,23 @@
+import argparse
+parser = argparse.ArgumentParser(description='SESR')
+
+# ipu test or cpu, you need to provide onnx path 
+parser.add_argument('--ipu', action='store_true',
+                    help='use ipu')
+parser.add_argument('--onnx_path', type=str, default='SESR_int8.onnx',
+                    help='onnx path')
+parser.add_argument('--provider_config', type=str, default=None,
+                    help='provider config path')
+# Data specifications, you can use default
+parser.add_argument('--dir_data', type=str, default='dataset/',
+                    help='dataset directory')
+parser.add_argument('--data_test', type=str, default='Set5',
+                    help='test dataset name')
+parser.add_argument('--n_threads', type=int, default=6,
+                    help='number of threads for data loading')
+parser.add_argument('--scale', type=str, default='2',
+                    help='super resolution scale')
+args = parser.parse_args()
+args.scale = list(map(lambda x: int(x), args.scale.split('+')))
+args.data_test = args.data_test.split('+')
+

requirements.txt

@@ -0,0 +1,8 @@
+torch==1.13.1
+numpy>=1.23.5
+scipy>=1.9
+opencv-python
+pandas
+pillow
+scikit-image
+tqdm

test.py

@@ -0,0 +1,59 @@
+import torch
+import sys
+import pathlib
+CURRENT_DIR = pathlib.Path(__file__).parent
+sys.path.append(str(CURRENT_DIR))
+from tqdm import tqdm
+import utility
+import data
+from option import args
+import metric
+import onnxruntime
+import cv2
+from data.data_tiling import tiling_inference
+
+
+def prepare(a, b, device):
+    def _prepare(tensor):
+        return tensor.to(device)
+
+    return _prepare(a), _prepare(b)
+
+
+def test_model(session, loader, device):
+    torch.set_grad_enabled(False)
+    self_scale = [2]
+    for idx_data, d in enumerate(loader.loader_test):
+        eval_ssim = 0
+        eval_psnr = 0
+        for idx_scale, scale in enumerate(self_scale):
+            d.dataset.set_scale(idx_scale)
+            for lr, hr, filename in tqdm(d, ncols=80):
+                lr, hr = prepare(lr, hr, device)
+                sr = tiling_inference(session, lr.cpu().numpy(), 8, (56, 56))
+                sr = torch.from_numpy(sr).to(device)
+                sr = utility.quantize(sr, 255)
+                eval_psnr += metric.calc_psnr(
+                    sr, hr, scale, 255, benchmark=d)
+                eval_ssim += metric.calc_ssim(
+                    sr, hr, scale, 255, dataset=d)
+            mean_ssim = eval_ssim / len(d)
+            mean_psnr = eval_psnr  / len(d)
+        print("psnr: %s, ssim: %s"%(mean_psnr, mean_ssim))
+    return mean_psnr, mean_ssim
+
+def main():
+    loader = data.Data(args)
+    if args.ipu:
+        providers = ["VitisAIExecutionProvider"]
+        provider_options = [{"config_file": args.provider_config}]
+    else:
+        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+        provider_options = None
+    onnx_file_name = args.onnx_path
+    ort_session = onnxruntime.InferenceSession(onnx_file_name, providers=providers, provider_options=provider_options) 
+    test_model(ort_session, loader, device="cpu")
+
+
+if __name__ == '__main__':
+    main()

utility.py

@@ -0,0 +1,33 @@
+import os
+import math
+import time
+import numpy as np
+
+class timer():
+    def __init__(self):
+        self.acc = 0
+        self.tic()
+
+    def tic(self):
+        self.t0 = time.time()
+
+    def toc(self, restart=False):
+        diff = time.time() - self.t0
+        if restart: self.t0 = time.time()
+        return diff
+
+    def hold(self):
+        self.acc += self.toc()
+
+    def release(self):
+        ret = self.acc
+        self.acc = 0
+
+        return ret
+
+    def reset(self):
+        self.acc = 0
+
+def quantize(img, rgb_range):
+    pixel_range = 255 / rgb_range
+    return img.mul(pixel_range).clamp(0, 255).round().div(pixel_range)