init model

RCAN_int8.onnx

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

README.md

@@ -0,0 +1,124 @@
+---
+license: apache-2.0
+tags:
+- RyzenAI
+- Super Resolution
+- Pytorch
+- Vision
+- SISR
+datasets:
+- Set5
+- Div2k
+language:
+- en
+Metircs:
+- PSNR
+---
+
+# RCAN model trained on DIV2K
+
+RCAN is a very deep residual channel attention network for super resolution trained on DIV2K. It was introduced in the paper [Image Super-Resolution Using Very Deep Residual Channel Attention Networks in 2018](https://arxiv.org/abs/1807.02758) by Yulun Zhang et al. and first released in [this repository](https://github.com/yulunzhang/RCAN).
+
+We develop a modified version that could be supported by [AMD Ryzen AI](https://ryzenai.docs.amd.com).
+
+
+## Model description
+RCAN is an advanced algorithm for single image super resolution. Our modified version is smaller than the original version. It is based deep learning techniques and is capable of X2 super resolution.
+
+
+## Intended uses & limitations
+
+You can use the raw model for super resolution. See the [model hub](https://huggingface.co/models?sort=trending&search=amd%2Frcan) to look for all available RCAN 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 [`infer_onnx.py`](infer_onnx.py) on how to use
+```python
+    parser = argparse.ArgumentParser(description='RCAN SISR')
+    parser.add_argument('--onnx_path', type=str, default='RCAN_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 infer_onnx.py --onnx_path RCAN_int8.onnx --image_path /Path/To/Your/Image --ipu --provider_config Path/To/vaip_config.json
+  ```
+
+ - Test accuracy of the quantized model
+  ```python
+  python eval_onnx.py --onnx_path RCAN_int8.onnx --data_test Set5 --ipu --provider_config Path/To/vaip_config.json 
+  ```
+### Performance
+
+| Method     | Scale | Flops | Set5         |
+|------------|-------|-------|--------------|
+|RCAN-S (float) |X2 |24.5G |37.531 / 0.958|
+|RCAN-S (INT8) |X2  |24.5G |37.150 / 0.955|
+- Note: the Flops is calculated with the output resolution is 360x640
+
+```bibtex
+@inproceedings{zhang2018image,
+  title={Image super-resolution using very deep residual channel attention networks},
+  author={Zhang, Yulun and Li, Kunpeng and Li, Kai and Wang, Lichen and Zhong, Bineng and Fu, Yun},
+  booktitle={Proceedings of the European conference on computer vision (ECCV)},
+  pages={286--301},
+  year={2018}
+}
+```

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,40 @@
+import os
+import onnxruntime
+import numpy as np
+import math
+
+
+def tiling_inference(session, lr, overlapping, patch_size):
+    _, _, 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,85 @@
+import os
+import glob
+import random
+import pickle
+
+from data import common
+
+import numpy as np
+import imageio
+import torch
+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)
+

eval_onnx.py

@@ -0,0 +1,87 @@
+import torch
+import os
+import sys
+import pathlib
+CURRENT_DIR = pathlib.Path(__file__).parent
+sys.path.append(str(CURRENT_DIR))
+from tqdm import tqdm
+import data
+import metric
+import onnxruntime
+import cv2
+from data.data_tiling import tiling_inference
+import argparse
+
+class Configs():
+    def __init__(self):
+        parser = argparse.ArgumentParser(description='SR')
+
+        # 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='RCAN_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, now only support x2')
+        self.parser = parser
+
+    def parse(self):
+        args = self.parser.parse_args()
+        args.scale = list(map(lambda x: int(x), args.scale.split('+')))
+        args.data_test = args.data_test.split('+')
+        print(args)
+        return args
+
+
+
+def quantize(img, rgb_range): # clamp pix to rgb range
+    pixel_range = 255 / rgb_range
+    return img.mul(pixel_range).clamp(0, 255).round().div(pixel_range)
+
+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):
+                sr = tiling_inference(session, lr.cpu().numpy(), 8, (56, 56))
+                sr = torch.from_numpy(sr).to(device)
+                sr = 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(args):
+    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__':
+    args = Configs().parse()
+    main(args)

infer_onnx.py

@@ -0,0 +1,45 @@
+import torch
+import onnxruntime
+import cv2
+import sys
+import pathlib
+CURRENT_DIR = pathlib.Path(__file__).parent
+sys.path.append(str(CURRENT_DIR))
+import numpy as np
+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)
+    sr = cv2.imwrite(output_path, sr)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='RCAN SISR')
+    parser.add_argument('--onnx_path', type=str, default='RCAN_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)

metric.py

@@ -0,0 +1,93 @@
+import math
+import numpy as np
+import torch
+import skimage.measure
+import skimage.color
+
+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)
+
+
+import numpy as np
+from scipy import signal
+
+
+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)
+
+  if True:#dataset and dataset.dataset.benchmark:
+    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)
+  else:
+    shave = scale + 6
+
+  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

requirements.txt

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