HairFastGAN

1
import os
2
import numpy as np
3
import torch
4
import torch.nn as nn
5
import torch.nn.functional as F
6
import torch.utils.data as data
7

8
from PIL import Image
9
from torch.autograd import grad
10

11
        
12
def clip_img(x):
13
    """Clip stylegan generated image to range(0,1)"""
14
    img_tmp = x.clone()[0]
15
    img_tmp = (img_tmp + 1) / 2
16
    img_tmp = torch.clamp(img_tmp, 0, 1)
17
    return [img_tmp.detach().cpu()]
18

19
def tensor_byte(x):
20
    return x.element_size()*x.nelement()
21

22
def count_parameters(net):
23
    s = sum([np.prod(list(mm.size())) for mm in net.parameters()])
24
    print(s)
25

26
def stylegan_to_classifier(x, out_size=(224, 224)):
27
    """Clip image to range(0,1)"""
28
    img_tmp = x.clone()
29
    img_tmp = torch.clamp((0.5*img_tmp + 0.5), 0, 1)
30
    img_tmp = F.interpolate(img_tmp, size=out_size, mode='bilinear')
31
    img_tmp[:,0] = (img_tmp[:,0] - 0.485)/0.229
32
    img_tmp[:,1] = (img_tmp[:,1] - 0.456)/0.224
33
    img_tmp[:,2] = (img_tmp[:,2] - 0.406)/0.225
34
    #img_tmp = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(img_tmp)
35
    return img_tmp
36
    
37
def downscale(x, scale_times=1, mode='bilinear'):
38
    for i in range(scale_times):
39
        x = F.interpolate(x, scale_factor=0.5, mode=mode)
40
    return x
41
    
42
def upscale(x, scale_times=1, mode='bilinear'):
43
    for i in range(scale_times):
44
        x = F.interpolate(x, scale_factor=2, mode=mode)
45
    return x
46
    
47
def hist_transform(source_tensor, target_tensor):
48
    """Histogram transformation"""
49
    c, h, w = source_tensor.size()
50
    s_t = source_tensor.view(c, -1)
51
    t_t = target_tensor.view(c, -1)
52
    s_t_sorted, s_t_indices = torch.sort(s_t)
53
    t_t_sorted, t_t_indices = torch.sort(t_t)
54
    for i in range(c):
55
        s_t[i, s_t_indices[i]] = t_t_sorted[i]
56
    return s_t.view(c, h, w)
57

58
def init_weights(m):
59
    """Initialize layers with Xavier uniform distribution"""
60
    if type(m) == nn.Conv2d:
61
        nn.init.xavier_uniform_(m.weight)
62
    elif type(m) == nn.Linear:
63
        nn.init.uniform_(m.weight, 0.0, 1.0)
64
        if m.bias is not None:
65
            nn.init.constant_(m.bias, 0.01)
66

67
def total_variation(x, delta=1):
68
    """Total variation, x: tensor of size (B, C, H, W)"""
69
    out = torch.mean(torch.abs(x[:, :, :, :-delta] - x[:, :, :, delta:]))\
70
        + torch.mean(torch.abs(x[:, :, :-delta, :] - x[:, :, delta:, :]))
71
    return out
72

73
def vgg_transform(x):
74
    """Adapt image for vgg network, x: image of range(0,1) subtracting ImageNet mean"""
75
    r, g, b = torch.split(x, 1, 1)
76
    out = torch.cat((b, g, r), dim = 1)
77
    out = F.interpolate(out, size=(224, 224), mode='bilinear')
78
    out = out*255.
79
    return out
80

81
# warp image with flow
82
def normalize_axis(x,L):
83
    return (x-1-(L-1)/2)*2/(L-1)
84

85
def unnormalize_axis(x,L):
86
    return x*(L-1)/2+1+(L-1)/2
87

88
def torch_flow_to_th_sampling_grid(flow,h_src,w_src,use_cuda=False):
89
    b,c,h_tgt,w_tgt=flow.size()
90
    grid_y, grid_x = torch.meshgrid(torch.tensor(range(1,w_tgt+1)),torch.tensor(range(1,h_tgt+1)))
91
    disp_x=flow[:,0,:,:]
92
    disp_y=flow[:,1,:,:]
93
    source_x=grid_x.unsqueeze(0).repeat(b,1,1).type_as(flow)+disp_x
94
    source_y=grid_y.unsqueeze(0).repeat(b,1,1).type_as(flow)+disp_y
95
    source_x_norm=normalize_axis(source_x,w_src) 
96
    source_y_norm=normalize_axis(source_y,h_src) 
97
    sampling_grid=torch.cat((source_x_norm.unsqueeze(3), source_y_norm.unsqueeze(3)), dim=3)
98
    if use_cuda:
99
        sampling_grid = sampling_grid.cuda()
100
    return sampling_grid
101

102
def warp_image_torch(image, flow):
103
    """
104
    Warp image (tensor, shape=[b, 3, h_src, w_src]) with flow (tensor, shape=[b, h_tgt, w_tgt, 2])
105
    """
106
    b,c,h_src,w_src=image.size()
107
    sampling_grid_torch = torch_flow_to_th_sampling_grid(flow, h_src, w_src)  
108
    warped_image_torch = F.grid_sample(image, sampling_grid_torch)
109
    return warped_image_torch