Update code and model to support NHWC input format

PAN_int8.onnx

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

data/benchmark.py

@@ -3,9 +3,9 @@ import os
 from data import srdata
 
 class Benchmark(srdata.SRData):
-    def __init__(self, args, name='',  benchmark=True):
+    def __init__(self, args, name='',  benchmark=True, input_data_format ='NHWC'):
         super(Benchmark, self).__init__(
-            args, name=name,  benchmark=True
+            args, name=name,  benchmark=True, input_data_format=input_data_format
         )
 
     def _set_filesystem(self, dir_data):

data/common.py

@@ -20,12 +20,13 @@ def set_channel(*args, n_channels=3):
 
     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()
+def np2Tensor(*args, rgb_range=255, format='NCHW'):
+    def _np2Tensor(img, channel_format):
+        assert channel_format in ('NCHW', 'NHWC')
+        img = np.ascontiguousarray(img.transpose((2, 0, 1))) if channel_format == ('NCHW') else img
+        tensor = torch.from_numpy(img).float()
         tensor.mul_(rgb_range / 255)
 
         return tensor
 
-    return [_np2Tensor(a) for a in args]
+    return [_np2Tensor(a, format) for a in args]

data/data_tiling.py

@@ -11,8 +11,8 @@ def tiling_inference(session, lr, overlapping=8, patch_size=(56, 56)):
         - 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))
+    _, h, w, _ = lr.shape
+    sr = np.zeros((1, 2*h, 2*w, 3))
     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
@@ -23,8 +23,9 @@ def tiling_inference(session, lr, overlapping=8, patch_size=(56, 56)):
             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]
+            tiling_lr = lr[..., h_idx: h_idx+patch_size[0], w_idx: w_idx+patch_size[1], :]
+            # import pdb; pdb.set_trace()
+            sr_tiling = session.run(None, {session.get_inputs()[0].name: tiling_lr})[0]
 
             left, right, top, bottom = 0, patch_size[1], 0, patch_size[0]
             left += overlapping//2 
@@ -42,5 +43,5 @@ def tiling_inference(session, lr, overlapping=8, patch_size=(56, 56)):
                 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]   
+            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

@@ -9,13 +9,15 @@ import imageio
 import torch.utils.data as data
 
 class SRData(data.Dataset):
-    def __init__(self, args, name='', benchmark=True):
+    def __init__(self, args, name='', benchmark=True, input_data_format='NCHW'):
         self.args = args
         self.name = name
         self.benchmark = benchmark
         self.input_large = False
         self.scale = args.scale
         self.idx_scale = 0
+        assert input_data_format in ('NCHW', 'NHWC')
+        self.input_data_format = input_data_format
         
         self._set_filesystem(args.dir_data)
         if args.ext.find('img') < 0:
@@ -87,7 +89,7 @@ class SRData(data.Dataset):
         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)
+        pair_t = common.np2Tensor(*pair, rgb_range=self.args.rgb_range, format=self.input_data_format)
 
         return pair_t[0], pair_t[1], filename
 

eval_onnx.py

@@ -26,6 +26,10 @@ def test_model(session, loader):
                 sr = tiling_inference(session, lr.numpy(), 8, (56, 56))
                 sr = torch.from_numpy(sr)
                 sr = utility.quantize(sr, 255)
+
+                # Transform from NHWC to NCHW to calculate metric
+                sr = sr.permute((0, 3, 1, 2))
+                hr = hr.permute((0, 3, 1, 2))
                 eval_psnr += utility.calc_psnr(
                     sr, hr, scale, 255, benchmark=d)
                 eval_ssim += utility.calc_ssim(

infer_onnx.py

@@ -22,12 +22,12 @@ def main(args):
         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)
+    lr = cv2.imread(image_path)[np.newaxis,:,:,:].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)
+    sr = sr.squeeze().astype(np.uint8)
     cv2.imwrite(output_path, sr)
 
 

utility.py

@@ -8,6 +8,9 @@ def quantize(img, 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(-1) == 3 and sr.size(1) > 3:
+       sr = sr.transpose((0, 3, 1, 2))
+       hr = hr.transpose((0, 3, 1, 2))
     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)]