init model

PAN_int8.onnx

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

README.md

@@ -0,0 +1,134 @@
+---
+license: apache-2.0
+datasets:
+- Set5
+- Div2K
+language:
+- en
+tags:
+- RyzenAI
+- PAN
+- Pytorch
+- Super Resolution
+- Vision
+pipeline_tag: image-to-image
+---
+
+## Model description
+
+PAN is an lightwight image super-resolution method with pixel pttention. It was introduced in the paper [Efficient Image Super-Resolution Using Pixel Attention](https://arxiv.org/abs/2010.01073) by Hengyuan Zhao et al. and first released in [this repository](https://github.com/zhaohengyuan1/PAN).
+
+We changed the negative slope of the leaky ReLU of the original model and replaced the sigmoid activation with hard sigmoid to make the model compatible with [AMD Ryzen AI](https://onnxruntime.ai/docs/execution-providers/Vitis-AI-ExecutionProvider.html). We loaded the published model parameters and fine-tuned them on the DIV2K dataset.
+
+
+## Intended uses & limitations
+
+You can use the raw model for super resolution. See the [model hub](https://huggingface.co/models?search=amd/pan) to look for all available PAN 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.
+3. Unzip the dataset and put it under the project folder. Organize the dataset directory as follows:
+```Plain
+PAN
+└── 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='PAN SR')
+    parser.add_argument('--onnx_path',
+                        type=str, 
+                        default='PAN_int8.onnx',
+                        help='Onnx path')
+    parser.add_argument('--image_path', 
+                        type=str,
+                        default='test_data/test.png',
+                        help='Path to your input image.')
+    parser.add_argument('--output_path', 
+                        type=str,
+                        default='test_data/sr.png',
+                        help='Path to your output image.')
+    parser.add_argument('--provider_config',
+                        type=str,
+                        default="vaip_config.json",
+                        help="Path of the config file for seting provider_options.")
+    parser.add_argument('--ipu', action='store_true', help='Use Ipu for interence.')
+
+    args = parser.parse_args()
+   
+    onnx_file_name = args.onnx_path
+    image_path = args.image_path
+    output_path = args.output_path
+
+    if args.ipu:
+        providers = ["VitisAIExecutionProvider"]
+        provider_options = [{"config_file": args.provider_config}]
+    else:
+        providers = ['CPUExecutionProvider']
+        provider_options = None
+    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 PAN_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 PAN_int8.onnx --data_test Set5 --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))
+
+### Performance
+
+| Method     | Scale | Flops | Set5         |
+|------------|-------|-------|--------------|
+|PAN (float) |X2 |141G |38.00 / 0.961|
+|PAN_amd (float) |X2  |141G |37.859 / 0.960|
+|PAN_amd (int8) |X2  |141G |37.18 / 0.952|
+- Note: the Flops is calculated with the output resolution is 360x640
+
+```bibtex
+@inproceedings{zhao2020efficient,
+  title={Efficient image super-resolution using pixel attention},
+  author={Zhao, Hengyuan and Kong, Xiangtao and He, Jingwen and Qiao, Yu and Dong, Chao},
+  booktitle={European Conference on Computer Vision},
+  pages={56--72},
+  year={2020},
+  organization={Springer}
+}
+```

data/__init__.py

@@ -0,0 +1,23 @@
+from importlib import import_module
+from torch.utils.data import dataloader
+
+
+class Data:
+    def __init__(self, args):
+        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:
+                raise 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,18 @@
+import os
+
+from data import srdata
+
+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,31 @@
+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=8, patch_size=(56, 56)):
+    """
+    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,130 @@
+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=True):
+        self.args = args
+        self.name = name
+        self.benchmark = benchmark
+        self.input_large = False
+        self.scale = args.scale
+        self.idx_scale = 0
+        
+        self._set_filesystem(args.dir_data)
+        if args.ext.find('img') < 0:
+            path_bin = os.path.join(self.apath, 'bin')
+            os.makedirs(path_bin, exist_ok=True)
+
+        list_hr, list_lr = self._scan()
+        if args.ext.find('img') >= 0 or benchmark:
+            self.images_hr, self.images_lr = list_hr, list_lr
+        elif args.ext.find('sep') >= 0:
+            os.makedirs(
+                self.dir_hr.replace(self.apath, path_bin),
+                exist_ok=True
+            )
+            for s in self.scale:
+                os.makedirs(
+                    os.path.join(
+                        self.dir_lr.replace(self.apath, path_bin),
+                        'X{}'.format(s)
+                    ),
+                    exist_ok=True
+                )
+            
+            self.images_hr, self.images_lr = [], [[] for _ in self.scale]
+            for h in list_hr:
+                b = h.replace(self.apath, path_bin)
+                b = b.replace(self.ext[0], '.pt')
+                self.images_hr.append(b)
+                self._check_and_load(args.ext, h, b, verbose=True) 
+            for i, ll in enumerate(list_lr):
+                for l in ll:
+                    b = l.replace(self.apath, path_bin)
+                    b = b.replace(self.ext[1], '.pt')
+                    self.images_lr[i].append(b)
+                    self._check_and_load(args.ext, l, b, verbose=True) 
+
+    # 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')
+        if self.input_large: self.dir_lr += 'L'
+        self.ext = ('.png', '.png')
+
+    def _check_and_load(self, ext, img, f, verbose=True):
+        if not os.path.isfile(f) or ext.find('reset') >= 0:
+            if verbose:
+                print('Making a binary: {}'.format(f))
+            with open(f, 'wb') as _f:
+                pickle.dump(imageio.imread(img), _f)
+
+    def __getitem__(self, idx):
+        lr, hr, filename = self._load_file(idx)
+        pair = self.get_patch(lr, hr)
+        pair = common.set_channel(*pair, n_channels=self.args.n_colors)
+        pair_t = common.np2Tensor(*pair, rgb_range=self.args.rgb_range)
+
+        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))
+        if self.args.ext == 'img' or self.benchmark:
+            hr = imageio.imread(f_hr)
+            lr = imageio.imread(f_lr)
+        elif self.args.ext.find('sep') >= 0:
+            with open(f_hr, 'rb') as _f:
+                hr = pickle.load(_f)
+            with open(f_lr, 'rb') as _f:
+                lr = pickle.load(_f)
+
+        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,52 @@
+import sys
+import pathlib
+CURRENT_DIR = pathlib.Path(__file__).parent
+sys.path.append(str(CURRENT_DIR))
+
+import torch
+from tqdm import tqdm
+import utility
+import data
+from option import args
+import onnxruntime
+from data.data_tiling import tiling_inference
+
+
+def test_model(session, loader):
+    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):
+
+                # Tiled inference
+                sr = tiling_inference(session, lr.numpy(), 8, (56, 56))
+                sr = torch.from_numpy(sr)
+                sr = utility.quantize(sr, 255)
+                eval_psnr += utility.calc_psnr(
+                    sr, hr, scale, 255, benchmark=d)
+                eval_ssim += utility.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)
+    onnx_file_name = args.onnx_path
+    if args.ipu:
+        providers = ["VitisAIExecutionProvider"]
+        provider_options = [{"config_file": args.provider_config}]
+    else:
+        providers = ['CPUExecutionProvider']
+        provider_options = None
+    ort_session = onnxruntime.InferenceSession(onnx_file_name,  providers=providers, provider_options=provider_options) 
+    test_model(ort_session, loader)
+
+
+if __name__ == '__main__':
+    main()

infer_onnx.py

@@ -0,0 +1,56 @@
+import sys
+import pathlib
+CURRENT_DIR = pathlib.Path(__file__).parent
+sys.path.append(str(CURRENT_DIR))
+
+import onnxruntime
+import cv2
+import numpy as np
+from data.data_tiling import tiling_inference
+import argparse
+
+
+def main(args):
+    onnx_file_name = args.onnx_path
+    image_path = args.image_path
+    output_path = args.output_path
+
+    if args.ipu:
+        providers = ["VitisAIExecutionProvider"]
+        provider_options = [{"config_file": args.provider_config}]
+    else:
+        providers = ['CPUExecutionProvider']
+        provider_options = None
+    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)
+
+    # Tiled inference
+    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='PAN')
+    parser.add_argument('--onnx_path',
+                        type=str, 
+                        default='PAN_int8.onnx',
+                        help='Path to onnx model')
+    parser.add_argument('--image_path', 
+                        type=str,
+                        default='test_data/test.png',
+                        help='Path to your low resolution input image.')
+    parser.add_argument('--output_path', 
+                        type=str,
+                        default='test_data/sr.png',
+                        help='Path to your upscaled output image.')
+    parser.add_argument('--provider_config',
+                        type=str,
+                        default="vaip_config.json",
+                        help="Path of the config file for seting provider_options.")
+    parser.add_argument('--ipu', action='store_true', help='Use Ipu for interence.')
+
+    args = parser.parse_args()
+    main(args)

option.py

@@ -0,0 +1,42 @@
+import argparse
+
+parser = argparse.ArgumentParser(description='PAN')
+
+# Hardware specifications
+parser.add_argument('--n_threads', type=int, default=6,
+                    help='Number of threads for data loading')
+parser.add_argument('--ipu', action='store_true', help='Use Ipu for interence.')
+
+# Data specifications
+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('--ext', type=str, default='sep',
+                    help='Dataset file extension')
+parser.add_argument('--scale', type=str, default='2',
+                    help='Super resolution scale')
+parser.add_argument('--rgb_range', type=int, default=255,
+                    help='Maximum value of RGB')
+parser.add_argument('--n_colors', type=int, default=3,
+                    help='Number of color channels to use')
+parser.add_argument('--onnx_path', type=str, default='PAN_int8.onnx',
+                    help='Path to onnx model')
+parser.add_argument('--provider_config',
+                    type=str,
+                    default="vaip_config.json",
+                    help="Path of the config file for seting provider_options.")
+
+
+args = parser.parse_args()
+
+args.scale = list(map(lambda x: int(x), args.scale.split('+')))
+args.data_test = args.data_test.split('+')
+
+
+for arg in vars(args):
+    if vars(args)[arg] == 'True':
+        vars(args)[arg] = True
+    elif vars(args)[arg] == 'False':
+        vars(args)[arg] = False
+

requirements.txt

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

utility.py

@@ -0,0 +1,89 @@
+import math
+import numpy as np
+from scipy import signal
+
+
+def quantize(img, rgb_range):
+    pixel_range = 255 / rgb_range
+    return img.mul(pixel_range).clamp(0, 255).round().div(pixel_range)
+
+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
+  '''
+  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
+
+