Spaces:
Runtime error
Runtime error
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| # | |
| # This source code is licensed under the BSD license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| 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: # gray image | |
| 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: | |
| # Unlike MATLAB, numpy.unit8() WILL NOT round by default. | |
| 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 | |