Spaces:
Runtime error
Runtime error
"""This module contains simple helper functions """ | |
from __future__ import print_function | |
import torch | |
import numbers | |
import torch.nn as nn | |
import torchvision | |
import torch.nn.functional as F | |
import math | |
import numpy as np | |
from PIL import Image | |
import os | |
import importlib | |
import argparse | |
from argparse import Namespace | |
from sklearn.decomposition import PCA as PCA | |
def normalize(v): | |
if type(v) == list: | |
return [normalize(vv) for vv in v] | |
return v * torch.rsqrt((torch.sum(v ** 2, dim=1, keepdim=True) + 1e-8)) | |
def slerp(a, b, r): | |
d = torch.sum(a * b, dim=-1, keepdim=True) | |
p = r * torch.acos(d * (1 - 1e-4)) | |
c = normalize(b - d * a) | |
d = a * torch.cos(p) + c * torch.sin(p) | |
return normalize(d) | |
def lerp(a, b, r): | |
if type(a) == list or type(a) == tuple: | |
return [lerp(aa, bb, r) for aa, bb in zip(a, b)] | |
return a * (1 - r) + b * r | |
def madd(a, b, r): | |
if type(a) == list or type(a) == tuple: | |
return [madd(aa, bb, r) for aa, bb in zip(a, b)] | |
return a + b * r | |
def str2bool(v): | |
if isinstance(v, bool): | |
return v | |
if v.lower() in ('yes', 'true', 't', 'y', '1'): | |
return True | |
elif v.lower() in ('no', 'false', 'f', 'n', '0'): | |
return False | |
else: | |
raise argparse.ArgumentTypeError('Boolean value expected.') | |
def copyconf(default_opt, **kwargs): | |
conf = Namespace(**vars(default_opt)) | |
for key in kwargs: | |
setattr(conf, key, kwargs[key]) | |
return conf | |
def find_class_in_module(target_cls_name, module): | |
target_cls_name = target_cls_name.replace('_', '').lower() | |
clslib = importlib.import_module(module) | |
cls = None | |
for name, clsobj in clslib.__dict__.items(): | |
if name.lower() == target_cls_name: | |
cls = clsobj | |
assert cls is not None, "In %s, there should be a class whose name matches %s in lowercase without underscore(_)" % (module, target_cls_name) | |
return cls | |
def tile_images(imgs, picturesPerRow=4): | |
""" Code borrowed from | |
https://stackoverflow.com/questions/26521365/cleanly-tile-numpy-array-of-images-stored-in-a-flattened-1d-format/26521997 | |
""" | |
# Padding | |
if imgs.shape[0] % picturesPerRow == 0: | |
rowPadding = 0 | |
else: | |
rowPadding = picturesPerRow - imgs.shape[0] % picturesPerRow | |
if rowPadding > 0: | |
imgs = np.concatenate([imgs, np.zeros((rowPadding, *imgs.shape[1:]), dtype=imgs.dtype)], axis=0) | |
# Tiling Loop (The conditionals are not necessary anymore) | |
tiled = [] | |
for i in range(0, imgs.shape[0], picturesPerRow): | |
tiled.append(np.concatenate([imgs[j] for j in range(i, i + picturesPerRow)], axis=1)) | |
tiled = np.concatenate(tiled, axis=0) | |
return tiled | |
# Converts a Tensor into a Numpy array | |
# |imtype|: the desired type of the converted numpy array | |
def tensor2im(image_tensor, imtype=np.uint8, normalize=True, tile=2): | |
if isinstance(image_tensor, list): | |
image_numpy = [] | |
for i in range(len(image_tensor)): | |
image_numpy.append(tensor2im(image_tensor[i], imtype, normalize)) | |
return image_numpy | |
if len(image_tensor.shape) == 4: | |
# transform each image in the batch | |
images_np = [] | |
for b in range(image_tensor.shape[0]): | |
one_image = image_tensor[b] | |
one_image_np = tensor2im(one_image) | |
images_np.append(one_image_np.reshape(1, *one_image_np.shape)) | |
images_np = np.concatenate(images_np, axis=0) | |
if tile is not False: | |
tile = max(min(images_np.shape[0] // 2, 4), 1) if tile is True else tile | |
images_tiled = tile_images(images_np, picturesPerRow=tile) | |
return images_tiled | |
else: | |
return images_np | |
if len(image_tensor.shape) == 2: | |
assert False | |
#imagce_tensor = image_tensor.unsqueeze(0) | |
image_numpy = image_tensor.detach().cpu().numpy() if type(image_tensor) is not np.ndarray else image_tensor | |
if normalize: | |
image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0 | |
else: | |
image_numpy = np.transpose(image_numpy, (1, 2, 0)) * 255.0 | |
image_numpy = np.clip(image_numpy, 0, 255) | |
if image_numpy.shape[2] == 1: | |
image_numpy = np.repeat(image_numpy, 3, axis=2) | |
return image_numpy.astype(imtype) | |
def toPILImage(images, tile=None): | |
if isinstance(images, list): | |
if all(['tensor' in str(type(image)).lower() for image in images]): | |
return toPILImage(torch.cat([im.cpu() for im in images], dim=0), tile) | |
return [toPILImage(image, tile=tile) for image in images] | |
if 'ndarray' in str(type(images)).lower(): | |
return toPILImage(torch.from_numpy(images)) | |
assert 'tensor' in str(type(images)).lower(), "input of type %s cannot be handled." % str(type(images)) | |
if tile is None: | |
max_width = 2560 | |
tile = min(images.size(0), int(max_width / images.size(3))) | |
return Image.fromarray(tensor2im(images, tile=tile)) | |
def diagnose_network(net, name='network'): | |
"""Calculate and print the mean of average absolute(gradients) | |
Parameters: | |
net (torch network) -- Torch network | |
name (str) -- the name of the network | |
""" | |
mean = 0.0 | |
count = 0 | |
for param in net.parameters(): | |
if param.grad is not None: | |
mean += torch.mean(torch.abs(param.grad.data)) | |
count += 1 | |
if count > 0: | |
mean = mean / count | |
print(name) | |
print(mean) | |
def save_image(image_numpy, image_path, aspect_ratio=1.0): | |
"""Save a numpy image to the disk | |
Parameters: | |
image_numpy (numpy array) -- input numpy array | |
image_path (str) -- the path of the image | |
""" | |
image_pil = Image.fromarray(image_numpy) | |
h, w, _ = image_numpy.shape | |
if aspect_ratio is None: | |
pass | |
elif aspect_ratio > 1.0: | |
image_pil = image_pil.resize((h, int(w * aspect_ratio)), Image.BICUBIC) | |
elif aspect_ratio < 1.0: | |
image_pil = image_pil.resize((int(h / aspect_ratio), w), Image.BICUBIC) | |
image_pil.save(image_path) | |
def print_numpy(x, val=True, shp=False): | |
"""Print the mean, min, max, median, std, and size of a numpy array | |
Parameters: | |
val (bool) -- if print the values of the numpy array | |
shp (bool) -- if print the shape of the numpy array | |
""" | |
x = x.astype(np.float64) | |
if shp: | |
print('shape,', x.shape) | |
if val: | |
x = x.flatten() | |
print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % ( | |
np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x))) | |
def mkdirs(paths): | |
"""create empty directories if they don't exist | |
Parameters: | |
paths (str list) -- a list of directory paths | |
""" | |
if isinstance(paths, list) and not isinstance(paths, str): | |
for path in paths: | |
mkdir(path) | |
else: | |
mkdir(paths) | |
def mkdir(path): | |
"""create a single empty directory if it didn't exist | |
Parameters: | |
path (str) -- a single directory path | |
""" | |
if not os.path.exists(path): | |
os.makedirs(path) | |
def visualize_spatial_code(sp): | |
device = sp.device | |
#sp = (sp - sp.min()) / (sp.max() - sp.min() + 1e-7) | |
if sp.size(1) <= 2: | |
sp = sp.repeat([1, 3, 1, 1])[:, :3, :, :] | |
if sp.size(1) == 3: | |
pass | |
else: | |
sp = sp.detach().cpu().numpy() | |
X = np.transpose(sp, (0, 2, 3, 1)) | |
B, H, W = X.shape[0], X.shape[1], X.shape[2] | |
X = np.reshape(X, (-1, X.shape[3])) | |
X = X - X.mean(axis=0, keepdims=True) | |
try: | |
Z = PCA(3).fit_transform(X) | |
except ValueError: | |
print("Running PCA on the structure code has failed.") | |
print("This is likely a bug of scikit-learn in version 0.18.1.") | |
print("https://stackoverflow.com/a/42764378") | |
print("The visualization of the structure code on visdom won't work.") | |
return torch.zeros(B, 3, H, W, device=device) | |
sp = np.transpose(np.reshape(Z, (B, H, W, -1)), (0, 3, 1, 2)) | |
sp = (sp - sp.min()) / (sp.max() - sp.min()) * 2 - 1 | |
sp = torch.from_numpy(sp).to(device) | |
return sp | |
def blank_tensor(w, h): | |
return torch.ones(1, 3, h, w) | |
class RandomSpatialTransformer: | |
def __init__(self, opt, bs): | |
self.opt = opt | |
#self.resample_transformation(bs) | |
def create_affine_transformation(self, ref, rot, sx, sy, tx, ty): | |
return torch.stack([-ref * sx * torch.cos(rot), -sy * torch.sin(rot), tx, | |
-ref * sx * torch.sin(rot), sy * torch.cos(rot), ty], axis=1) | |
def resample_transformation(self, bs, device, reflection=None, rotation=None, scale=None, translation=None): | |
dev = device | |
zero = torch.zeros((bs), device=dev) | |
if reflection is None: | |
#if "ref" in self.opt.random_transformation_mode: | |
ref = torch.round(torch.rand((bs), device=dev)) * 2 - 1 | |
#else: | |
# ref = 1.0 | |
else: | |
ref = reflection | |
if rotation is None: | |
#if "rot" in self.opt.random_transformation_mode: | |
max_rotation = 30 * math.pi / 180 | |
rot = torch.rand((bs), device=dev) * (2 * max_rotation) - max_rotation | |
#else: | |
# rot = 0.0 | |
else: | |
rot = rotation | |
if scale is None: | |
#if "scale" in self.opt.random_transformation_mode: | |
min_scale = 1.0 | |
max_scale = 1.0 | |
sx = torch.rand((bs), device=dev) * (max_scale - min_scale) + min_scale | |
sy = torch.rand((bs), device=dev) * (max_scale - min_scale) + min_scale | |
#else: | |
# sx, sy = 1.0, 1.0 | |
else: | |
sx, sy = scale | |
tx, ty = zero, zero | |
A = torch.stack([ref * sx * torch.cos(rot), -sy * torch.sin(rot), tx, | |
ref * sx * torch.sin(rot), sy * torch.cos(rot), ty], axis=1) | |
return A.view(bs, 2, 3) | |
def forward_transform(self, x, size): | |
if type(x) == list: | |
return [self.forward_transform(xx) for xx in x] | |
affine_param = self.resample_transformation(x.size(0), x.device) | |
affine_grid = F.affine_grid(affine_param, (x.size(0), x.size(1), size[0], size[1]), align_corners=False) | |
x = F.grid_sample(x, affine_grid, padding_mode='reflection', align_corners=False) | |
return x | |
def apply_random_crop(x, target_size, scale_range, num_crops=1, return_rect=False): | |
# build grid | |
B = x.size(0) * num_crops | |
flip = torch.round(torch.rand(B, 1, 1, 1, device=x.device)) * 2 - 1.0 | |
unit_grid_x = torch.linspace(-1.0, 1.0, target_size, device=x.device)[np.newaxis, np.newaxis, :, np.newaxis].repeat(B, target_size, 1, 1) | |
unit_grid_y = unit_grid_x.transpose(1, 2) | |
unit_grid = torch.cat([unit_grid_x * flip, unit_grid_y], dim=3) | |
#crops = [] | |
x = x.unsqueeze(1).expand(-1, num_crops, -1, -1, -1).flatten(0, 1) | |
#for i in range(num_crops): | |
scale = torch.rand(B, 1, 1, 2, device=x.device) * (scale_range[1] - scale_range[0]) + scale_range[0] | |
offset = (torch.rand(B, 1, 1, 2, device=x.device) * 2 - 1) * (1 - scale) | |
sampling_grid = unit_grid * scale + offset | |
crop = F.grid_sample(x, sampling_grid, align_corners=False) | |
#crops.append(crop) | |
#crop = torch.stack(crops, dim=1) | |
crop = crop.view(B // num_crops, num_crops, crop.size(1), crop.size(2), crop.size(3)) | |
return crop | |
def five_crop_noresize(A): | |
Y, X = A.size(2) // 3, A.size(3) // 3 | |
H, W = Y * 2, X * 2 | |
return torch.stack([A[:, :, 0:0+H, 0:0+W], | |
A[:, :, Y:Y+H, 0:0+W], | |
A[:, :, Y:Y+H, X:X+W], | |
A[:, :, 0:0+H, X:X+W], | |
A[:, :, Y//2:Y//2+H, X//2:X//2+W]], | |
dim=1) # return 5-dim tensor | |
def random_crop_noresize(A, crop_size): | |
offset_y = np.random.randint(A.size(2) - crop_size[0]) | |
offset_x = np.random.randint(A.size(3) - crop_size[1]) | |
return A[:, :, offset_y:offset_y + crop_size[0], offset_x:offset_x + crop_size[1]], (offset_y, offset_x) | |
def random_crop_with_resize(A, crop_size): | |
#size_y = np.random.randint(crop_size[0], A.size(2) + 1) | |
#size_x = np.random.randint(crop_size[1], A.size(3) + 1) | |
#size_y, size_x = crop_size | |
size_y = max(crop_size[0], np.random.randint(A.size(2) // 3, A.size(2) + 1)) | |
size_x = max(crop_size[1], np.random.randint(A.size(3) // 3, A.size(3) + 1)) | |
offset_y = np.random.randint(A.size(2) - size_y + 1) | |
offset_x = np.random.randint(A.size(3) - size_x + 1) | |
crop_rect = (offset_y, offset_x, size_y, size_x) | |
resized = crop_with_resize(A, crop_rect, crop_size) | |
#print('resized %s to %s' % (A.size(), resized.size())) | |
return resized, crop_rect | |
def crop_with_resize(A, crop_rect, return_size): | |
offset_y, offset_x, size_y, size_x = crop_rect | |
crop = A[:, :, offset_y:offset_y + size_y, offset_x:offset_x + size_x] | |
resized = F.interpolate(crop, size=return_size, mode='bilinear', align_corners=False) | |
#print('resized %s to %s' % (A.size(), resized.size())) | |
return resized | |
def compute_similarity_logit(x, y, p=1, compute_interdistances=True): | |
def compute_dist(x, y, p): | |
if p == 2: | |
return ((x - y) ** 2).sum(dim=-1).sqrt() | |
else: | |
return (x - y).abs().sum(dim=-1) | |
C = x.shape[-1] | |
if len(x.shape) == 2: | |
if compute_interdistances: | |
dist = torch.cdist(x[None, :, :], y[None, :, :], p)[0] | |
else: | |
dist = compute_dist(x, y, p) | |
if len(x.shape) == 3: | |
if compute_interdistances: | |
dist = torch.cdist(x, y, p) | |
else: | |
dist = compute_dist(x, y, p) | |
if p == 1: | |
dist = 1 - dist / math.sqrt(C) | |
elif p == 2: | |
dist = 1 - 0.5 * (dist ** 2) | |
return dist / 0.07 | |
def set_diag_(x, value): | |
assert x.size(-2) == x.size(-1) | |
L = x.size(-2) | |
identity = torch.eye(L, dtype=torch.bool, device=x.device) | |
identity = identity.view([1] * (len(x.shape) - 2) + [L, L]) | |
x.masked_fill_(identity, value) | |
def to_numpy(metric_dict): | |
new_dict = {} | |
for k, v in metric_dict.items(): | |
if "numpy" not in str(type(v)): | |
v = v.detach().cpu().mean().numpy() | |
new_dict[k] = v | |
return new_dict | |
def is_custom_kernel_supported(): | |
version_str = str(torch.version.cuda).split(".") | |
major = version_str[0] | |
minor = version_str[1] | |
return int(major) >= 10 and int(minor) >= 1 | |
def shuffle_batch(x): | |
B = x.size(0) | |
perm = torch.randperm(B, dtype=torch.long, device=x.device) | |
return x[perm] | |
def unravel_index(index, shape): | |
out = [] | |
for dim in reversed(shape): | |
out.append(index % dim) | |
index = index // dim | |
return tuple(reversed(out)) | |
def quantize_color(x, num=64): | |
return (x * num / 2).round() * (2 / num) | |
def resize2d_tensor(x, size_or_tensor_of_size): | |
if torch.is_tensor(size_or_tensor_of_size): | |
size = size_or_tensor_of_size.size() | |
elif isinstance(size_or_tensor_of_size, np.ndarray): | |
size = size_or_tensor_of_size.shape | |
else: | |
size = size_or_tensor_of_size | |
if isinstance(size, tuple) or isinstance(size, list): | |
return F.interpolate(x, size[-2:], | |
mode='bilinear', align_corners=False) | |
else: | |
raise ValueError("%s is unrecognized" % str(type(size))) | |
def correct_resize(t, size, mode=Image.BICUBIC): | |
device = t.device | |
t = t.detach().cpu() | |
resized = [] | |
for i in range(t.size(0)): | |
one_t = t[i:i+1] | |
one_image = Image.fromarray(tensor2im(one_t, tile=1)).resize(size, Image.BICUBIC) | |
resized_t = torchvision.transforms.functional.to_tensor(one_image) * 2 - 1.0 | |
resized.append(resized_t) | |
return torch.stack(resized, dim=0).to(device) | |
class GaussianSmoothing(nn.Module): | |
""" | |
Apply gaussian smoothing on a | |
1d, 2d or 3d tensor. Filtering is performed seperately for each channel | |
in the input using a depthwise convolution. | |
Arguments: | |
channels (int, sequence): Number of channels of the input tensors. Output will | |
have this number of channels as well. | |
kernel_size (int, sequence): Size of the gaussian kernel. | |
sigma (float, sequence): Standard deviation of the gaussian kernel. | |
dim (int, optional): The number of dimensions of the data. | |
Default value is 2 (spatial). | |
""" | |
def __init__(self, channels, kernel_size, sigma, dim=2): | |
super(GaussianSmoothing, self).__init__() | |
if isinstance(kernel_size, numbers.Number): | |
self.pad_size = kernel_size // 2 | |
kernel_size = [kernel_size] * dim | |
else: | |
raise NotImplementedError() | |
if isinstance(sigma, numbers.Number): | |
sigma = [sigma] * dim | |
# The gaussian kernel is the product of the | |
# gaussian function of each dimension. | |
kernel = 1 | |
meshgrids = torch.meshgrid( | |
[ | |
torch.arange(size, dtype=torch.float32) | |
for size in kernel_size | |
] | |
) | |
for size, std, mgrid in zip(kernel_size, sigma, meshgrids): | |
mean = (size - 1) / 2 | |
kernel *= 1 / (std * math.sqrt(2 * math.pi)) * \ | |
torch.exp(-((mgrid - mean) / std) ** 2 / 2) | |
# Make sure sum of values in gaussian kernel equals 1. | |
kernel = kernel / (torch.sum(kernel)) | |
# Reshape to depthwise convolutional weight | |
kernel = kernel.view(1, 1, *kernel.size()) | |
kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1)) | |
self.register_buffer('weight', kernel) | |
self.groups = channels | |
if dim == 1: | |
self.conv = F.conv1d | |
elif dim == 2: | |
self.conv = F.conv2d | |
elif dim == 3: | |
self.conv = F.conv3d | |
else: | |
raise RuntimeError( | |
'Only 1, 2 and 3 dimensions are supported. Received {}.'.format(dim) | |
) | |
def forward(self, input): | |
""" | |
Apply gaussian filter to input. | |
Arguments: | |
input (torch.Tensor): Input to apply gaussian filter on. | |
Returns: | |
filtered (torch.Tensor): Filtered output. | |
""" | |
x = F.pad(input, [self.pad_size] * 4, mode="reflect") | |
return self.conv(x, weight=self.weight, groups=self.groups) | |