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()
|