|
|
|
|
|
|
|
|
|
|
|
|
|
import os |
|
import cv2 |
|
import numpy as np |
|
import torch |
|
from os import path as osp |
|
from torch.nn import functional as F |
|
|
|
|
|
def scandir(dir_path, suffix=None, recursive=False, full_path=False): |
|
"""Scan a directory to find the interested files. |
|
|
|
Args: |
|
dir_path (str): Path of the directory. |
|
suffix (str | tuple(str), optional): File suffix that we are |
|
interested in. Default: None. |
|
recursive (bool, optional): If set to True, recursively scan the |
|
directory. Default: False. |
|
full_path (bool, optional): If set to True, include the dir_path. |
|
Default: False. |
|
|
|
Returns: |
|
A generator for all the interested files with relative paths. |
|
""" |
|
|
|
if (suffix is not None) and not isinstance(suffix, (str, tuple)): |
|
raise TypeError('"suffix" must be a string or tuple of strings') |
|
|
|
root = dir_path |
|
|
|
def _scandir(dir_path, suffix, recursive): |
|
for entry in os.scandir(dir_path): |
|
if not entry.name.startswith('.') and entry.is_file(): |
|
if full_path: |
|
return_path = entry.path |
|
else: |
|
return_path = osp.relpath(entry.path, root) |
|
|
|
if suffix is None: |
|
yield return_path |
|
elif return_path.endswith(suffix): |
|
yield return_path |
|
else: |
|
if recursive: |
|
yield from _scandir(entry.path, suffix=suffix, recursive=recursive) |
|
else: |
|
continue |
|
|
|
return _scandir(dir_path, suffix=suffix, recursive=recursive) |
|
|
|
|
|
def read_img_seq(path, require_mod_crop=False, scale=1, return_imgname=False): |
|
"""Read a sequence of images from a given folder path. |
|
|
|
Args: |
|
path (list[str] | str): List of image paths or image folder path. |
|
require_mod_crop (bool): Require mod crop for each image. |
|
Default: False. |
|
scale (int): Scale factor for mod_crop. Default: 1. |
|
return_imgname(bool): Whether return image names. Default False. |
|
|
|
Returns: |
|
Tensor: size (t, c, h, w), RGB, [0, 1]. |
|
list[str]: Returned image name list. |
|
""" |
|
if isinstance(path, list): |
|
img_paths = path |
|
else: |
|
img_paths = sorted(list(scandir(path, full_path=True))) |
|
imgs = [cv2.imread(v).astype(np.float32) / 255. for v in img_paths] |
|
|
|
if require_mod_crop: |
|
imgs = [mod_crop(img, scale) for img in imgs] |
|
imgs = img2tensor(imgs, bgr2rgb=True, float32=True) |
|
imgs = torch.stack(imgs, dim=0) |
|
|
|
if return_imgname: |
|
imgnames = [osp.splitext(osp.basename(path))[0] for path in img_paths] |
|
return imgs, imgnames |
|
else: |
|
return imgs |
|
|
|
|
|
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: |
|
if img.dtype == 'float64': |
|
img = img.astype('float32') |
|
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 tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)): |
|
"""Convert torch Tensors into image numpy arrays. |
|
|
|
After clamping to [min, max], values will be normalized to [0, 1]. |
|
|
|
Args: |
|
tensor (Tensor or list[Tensor]): Accept shapes: |
|
1) 4D mini-batch Tensor of shape (B x 3/1 x H x W); |
|
2) 3D Tensor of shape (3/1 x H x W); |
|
3) 2D Tensor of shape (H x W). |
|
Tensor channel should be in RGB order. |
|
rgb2bgr (bool): Whether to change rgb to bgr. |
|
out_type (numpy type): output types. If ``np.uint8``, transform outputs |
|
to uint8 type with range [0, 255]; otherwise, float type with |
|
range [0, 1]. Default: ``np.uint8``. |
|
min_max (tuple[int]): min and max values for clamp. |
|
|
|
Returns: |
|
(Tensor or list): 3D ndarray of shape (H x W x C) OR 2D ndarray of |
|
shape (H x W). The channel order is BGR. |
|
""" |
|
if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))): |
|
raise TypeError(f'tensor or list of tensors expected, got {type(tensor)}') |
|
|
|
if torch.is_tensor(tensor): |
|
tensor = [tensor] |
|
result = [] |
|
for _tensor in tensor: |
|
_tensor = _tensor.squeeze(0).float().detach().cpu().clamp_(*min_max) |
|
_tensor = (_tensor - min_max[0]) / (min_max[1] - min_max[0]) |
|
|
|
n_dim = _tensor.dim() |
|
if n_dim == 4: |
|
img_np = make_grid(_tensor, nrow=int(math.sqrt(_tensor.size(0))), normalize=False).numpy() |
|
img_np = img_np.transpose(1, 2, 0) |
|
if rgb2bgr: |
|
img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) |
|
elif n_dim == 3: |
|
img_np = _tensor.numpy() |
|
img_np = img_np.transpose(1, 2, 0) |
|
if img_np.shape[2] == 1: |
|
img_np = np.squeeze(img_np, axis=2) |
|
else: |
|
if rgb2bgr: |
|
img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) |
|
elif n_dim == 2: |
|
img_np = _tensor.numpy() |
|
else: |
|
raise TypeError(f'Only support 4D, 3D or 2D tensor. But received with dimension: {n_dim}') |
|
if out_type == np.uint8: |
|
|
|
img_np = (img_np * 255.0).round() |
|
img_np = img_np.astype(out_type) |
|
result.append(img_np) |
|
if len(result) == 1: |
|
result = result[0] |
|
return result |
|
|