|
- # ------------------------------------------------------------------------
- # Modified from BasicSR (https://github.com/xinntao/BasicSR)
- # Copyright 2018-2020 BasicSR Authors
- # ------------------------------------------------------------------------
-
- import cv2
- import torch
-
-
- def img2tensor(imgs, bgr2rgb=True, float32=True):
- """Numpy array to tensor.
- Args:
- imgs (list[ndarray] | ndarray): Input images.
- bgr2rgb (bool): Whether to change bgr to rgb.
- float32 (bool): Whether to change to float32.
- Returns:
- list[tensor] | tensor: Tensor images. If returned results only have
- one element, just return tensor.
- """
-
- def _totensor(img, bgr2rgb, float32):
- if img.shape[2] == 3 and bgr2rgb:
- img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
- img = torch.from_numpy(img.transpose(2, 0, 1))
- if float32:
- img = img.float()
- return img
-
- if isinstance(imgs, list):
- return [_totensor(img, bgr2rgb, float32) for img in imgs]
- else:
- return _totensor(imgs, bgr2rgb, float32)
-
-
- def padding(img_lq, img_gt, gt_size):
- h, w, _ = img_lq.shape
-
- h_pad = max(0, gt_size - h)
- w_pad = max(0, gt_size - w)
-
- if h_pad == 0 and w_pad == 0:
- return img_lq, img_gt
-
- img_lq = cv2.copyMakeBorder(img_lq, 0, h_pad, 0, w_pad, cv2.BORDER_REFLECT)
- img_gt = cv2.copyMakeBorder(img_gt, 0, h_pad, 0, w_pad, cv2.BORDER_REFLECT)
- return img_lq, img_gt
|
