|
import math
|
|
import yaml
|
|
import torch.nn.init as init
|
|
import torch
|
|
import numpy as np
|
|
|
|
def get_config(config):
|
|
with open(config, 'r') as stream:
|
|
return yaml.load(stream, Loader=yaml.Loader)
|
|
|
|
def weights_init(init_type='gaussian'):
|
|
def init_fun(m):
|
|
classname = m.__class__.__name__
|
|
if (classname.find('Conv') == 0 or classname.find('Linear') == 0) and hasattr(m, 'weight'):
|
|
|
|
if init_type == 'gaussian':
|
|
init.normal_(m.weight.data, 0.0, 0.02)
|
|
elif init_type == 'xavier':
|
|
init.xavier_normal_(m.weight.data, gain=math.sqrt(2))
|
|
elif init_type == 'kaiming':
|
|
init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
|
|
elif init_type == 'orthogonal':
|
|
init.orthogonal_(m.weight.data, gain=math.sqrt(2))
|
|
elif init_type == 'default':
|
|
pass
|
|
else:
|
|
assert 0, "Unsupported initialization: {}".format(init_type)
|
|
if hasattr(m, 'bias') and m.bias is not None:
|
|
init.constant_(m.bias.data, 0.0)
|
|
|
|
return init_fun
|
|
|
|
def tensor2im(input_image, imtype=np.uint8, no_fg=True):
|
|
""""Converts a Tensor array into a numpy image array.
|
|
|
|
Parameters:
|
|
input_image (tensor) -- the input image tensor array
|
|
imtype (type) -- the desired type of the converted numpy array
|
|
no_fg: binary image and don't transform
|
|
"""
|
|
if not isinstance(input_image, np.ndarray):
|
|
if isinstance(input_image, torch.Tensor):
|
|
image_tensor = input_image.data
|
|
else:
|
|
return input_image
|
|
image_numpy = image_tensor[0].cpu().float().numpy()
|
|
if no_fg:
|
|
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)
|
|
else:
|
|
image_numpy = input_image
|
|
return image_numpy.astype(imtype)
|
|
|