File size: 1,416 Bytes
96eb931
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""
usage:
python3 generate.py --image_path ./apple_test.jpg --name apple2orange --model cycle_gan --gpu_ids -1

gpu_ids: -1 for cpu inference
"""
from options.test_options import TestOptions
opt = TestOptions().parse()
from models.one_direction_test_model import OneDirectionTestModel
from data.unaligned_data_loader import load_image_for_prediction
import sys
import cv2
import os
import numpy as np
from PIL import Image

opt.nThreads = 1
opt.batchSize = 1
opt.serial_batches = True


def generate():
    """
    generate single image specific by image path, and show the after generated image
    :return:
    """
    image_path = opt.image_path
    print('generate from {}'.format(image_path))

    data = load_image_for_prediction(opt, image_path)

    model = OneDirectionTestModel()
    model.initialize(opt=opt)
    model.set_input(data)
    model.test()

    visuals = model.get_current_visuals()
    generated_a = visuals['fake_B']

    image_generated = Image.fromarray(generated_a)
    image_generated.save(str(os.path.basename(image_path).split('.')[0]) + '_fake_b.jpg')

    combined_result = np.concatenate([img for _, img in visuals.items()], 1)
    image_combined = Image.fromarray(combined_result)
    image_combined.save(str(os.path.basename(image_path).split('.')[0]) + '_combined.jpg')
    image_combined.show()
    print('generated image saved.')



if __name__ == '__main__':
    generate()