automodulator/eval_example.py at master · heljakka/automodulator · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
import torch
from PIL import Image
import numpy as np
import torchvision.utils

# Use as:
# z_fused = model.zbuilder().use(zA,[0,1]).use(zB,[2,3]).use(zA,[3,-1])
# or:
# z_fused = model.zbuilder().high(zA).mid(zB).lo(zA)
# imgAB = model.decode(z_fused)

repo = 'heljakka/automodulator'

#datatype = 'celebahq256'
#datatype = 'ffhq256'
datatype = 'ffhq512'
#datatype = 'lsuncars256'
#datatype = 'lsunbedrooms256'

model = torch.hub.load(repo, datatype, pretrained=True, source='github', force_reload = True)
model.eval(useLN=False)

# If you do not need the Low-memory footprint approach, just run model.generate(N) once for the N you need.
for k in range(4):
    omg = model.generate(1)
    torchvision.utils.save_image(omg[0]/ 2 + 0.5, fp=f'{datatype}_rand_{k}.png')
    del omg

if datatype in ['celebahq256', 'ffhq256', 'ffhq512']:
    simg = {}
    simg[0] = 'fig/source-0.png' #'/media/ari/ExData/data/ffhq_valid/ffhq_valid/60013.png'
    simg[1] = 'fig/source-1.png' #'/media/ari/ExData/data/ffhq_valid/ffhq_valid/60061.png'

    # Load 2 images to torch
    imgs = torch.stack([model.tf()(Image.open(simg[0])),
                        model.tf()(Image.open(simg[1]))])

    z = model.encode(imgs)

    # Reconstruct
    for i,img in enumerate(simg):
        omg = model.decode( model.zbuilder().hi(z[i]).mid(z[i]).lo(z[i]))
        torchvision.utils.save_image(omg[0], normalize=True, fp=f'{datatype}_reco_{i}.png')

    # Create typical mixture images
    omg = model.decode( model.zbuilder().hi(z[0]).mid(z[1]).lo(z[1]))
    torchvision.utils.save_image(omg[0], normalize=True, fp=f'{datatype}_mixB_011.png')

    omg = model.decode( model.zbuilder().hi(z[1]).mid(z[0]).lo(z[1]) )
    torchvision.utils.save_image(omg[0], normalize=True, fp=f'{datatype}_mixB_101.png')

    omg = model.decode( model.zbuilder().hi(z[1]).mid(z[1]).lo(z[0]) )
    torchvision.utils.save_image(omg[0], normalize=True, fp=f'{datatype}_mixB_110.png')

    omg = model.decode( model.zbuilder().hi(z[0]).mid(z[1]).lo(z[0]) )
    torchvision.utils.save_image(omg[0], normalize=True, fp=f'{datatype}_mixB_010.png')

    omg = model.decode( model.zbuilder().hi(z[1]).mid(z[0]).lo(z[0]) )
    torchvision.utils.save_image(omg[0], normalize=True, fp=f'{datatype}_mixB_100.png')

    # Mix two latents at every possible cut-off point (you could also interleave, not shown here)

    omgs = torch.stack([model.decode( model.zbuilder().use(z[0], [0,i]).use(z[1], [i,-1]) )
                        for i in range(0,9)])

    for j in range(9):
        torchvision.utils.save_image(omgs[j], normalize=True, fp=f'mix_{j}.png')