Spaces:
Sleeping
Sleeping
from enum import Enum | |
import numpy as np | |
from PIL import Image | |
import yaml | |
from easydict import EasyDict as edict | |
import torch.nn as nn | |
import torch | |
import cv2 | |
def load_image(fname, mode='RGB', return_orig=False): | |
img = np.array(Image.open(fname).convert(mode)) | |
if img.ndim == 3: | |
img = np.transpose(img, (2, 0, 1)) | |
out_img = img.astype('float32') / 255 | |
if return_orig: | |
return out_img, img | |
else: | |
return out_img | |
def prepare_image(input_img: Image, mode='RGB', return_orig=False): | |
img = np.array(input_img.convert(mode)) | |
if img.ndim == 3: | |
img = np.transpose(img, (2, 0, 1)) | |
out_img = img.astype('float32') / 255 | |
if return_orig: | |
return out_img, img | |
else: | |
return out_img | |
def ceil_modulo(x, mod): | |
if x % mod == 0: | |
return x | |
return (x // mod + 1) * mod | |
def pad_img_to_modulo(img, mod): | |
channels, height, width = img.shape | |
out_height = ceil_modulo(height, mod) | |
out_width = ceil_modulo(width, mod) | |
return np.pad(img, ((0, 0), (0, out_height - height), (0, out_width - width)), mode='symmetric') | |
def scale_image(img, factor, interpolation=cv2.INTER_AREA): | |
if img.shape[0] == 1: | |
img = img[0] | |
else: | |
img = np.transpose(img, (1, 2, 0)) | |
img = cv2.resize(img, dsize=None, fx=factor, fy=factor, interpolation=interpolation) | |
if img.ndim == 2: | |
img = img[None, ...] | |
else: | |
img = np.transpose(img, (2, 0, 1)) | |
return img | |
def load_yaml(path): | |
with open(path, 'r') as f: | |
return edict(yaml.safe_load(f)) | |
def move_to_device(obj, device): | |
if isinstance(obj, nn.Module): | |
return obj.to(device) | |
if torch.is_tensor(obj): | |
return obj.to(device) | |
if isinstance(obj, (tuple, list)): | |
return [move_to_device(el, device) for el in obj] | |
if isinstance(obj, dict): | |
return {name: move_to_device(val, device) for name, val in obj.items()} | |
raise ValueError(f'Unexpected type {type(obj)}') | |
class SmallMode(Enum): | |
DROP = "drop" | |
UPSCALE = "upscale" | |