ContextEncoder/test.lua at master · afrozalm/ContextEncoder · 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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
require 'image';
require 'torch';
require 'nn';
require 'cunn';
require 'cutorch';
require 'paths';

cmd = torch.CmdLine()
cmd:option('-old', 0)
cmd:option('-device', 1)
cmd:option('-start', 1)
cmd:option('-stop', 1)
params = cmd:parse(arg)

cutorch.setDevice(params['device'])
cutorch.synchronize()

generator =  nil
discriminator =  nil
classifier = nil

print('Loading Models ...')

if params['old'] == 1 then
    generatorName = './generatorOld.net'
    discriminatorName = './discriminatorOld.net'
else
    generatorName = './generator.net'
    discriminatorName = './discriminator.net'
end

if path.exists(generatorName) then
    generator = torch.load(generatorName)
    print('generator loaded')
else
    print('generator not available')
end
if path.exists(discriminatorName) then
    discriminator = torch.load(discriminatorName)
    print('discriminator loaded')
else
    print('discriminator not available')
end

print('Loading Data ... ')
testset = torch.load('../FaceScrub/FaceScrub_testset_128x128')

function redoutline( image )
    image[{{1},{32, 97},{32}}] = 255
    image[{{2, 3},{32, 97},{32}}] = 0
    image[{{1},{32, 97},{97}}] = 255
    image[{{2, 3},{32, 97},{97}}] = 0
    image[{{1},{32},{32, 97}}] = 255
    image[{{2, 3},{32},{32, 97}}] = 0
    image[{{1},{97},{32, 97}}] = 255
    image[{{2, 3},{97},{32, 97}}] = 0
    return image
end

M = torch.CudaTensor( 3, 128, 128 ):fill(1)      -- filling mask
M[{ {}, {33, 96}, {33, 96} }] = 0
print('Loaded')
print('Processing ...')

for index = params['start'], params['stop'] do
    fullInput = torch.CudaTensor( 3, 128, 128 ):copy( testset.data[index] )
    maskedInput = torch.CudaTensor( 3, 128, 128 ):copy( testset.data[index] ):cmul( M )
    if generator ~= nil then
        local temp4DTensor = torch.CudaTensor( 2, 3, 128, 128 )
        temp4DTensor[1] = maskedInput
        temp4DTensor[2] = fullInput
        output = generator:forward( temp4DTensor )
        outputImage = output[1]
        maskedInput[{ {}, {33, 96}, {33, 96} }] = outputImage
        temp4DTensor[1] = maskedInput
        probability = discriminator:forward(temp4DTensor)
        print( 'The fake is real with probability ' , probability[1] )
        print( 'Real is real with probability', probability[2] )
        print('Saving results')
        outname = index .. '_output.jpg'
        inname = index .. '_input.jpg'
        image.save( outname, redoutline( maskedInput ) )
        image.save( inname, redoutline( fullInput ) )
        temp4DTensor = nil
        collectgarbage()
    end
end