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