Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,339 Bytes
77771e4 |
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 |
import os
from options.test_options import TestOptions
from data import create_dataset
from models import create_model
from util.visualizer import save_images
from itertools import islice
from util import html
import cv2
seed = 10
import torch
import numpy as np
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
np.random.seed(seed)
# options
opt = TestOptions().parse()
opt.num_threads = 1 # test code only supports num_threads=1
opt.batch_size = 1 # test code only supports batch_size=1
opt.serial_batches = True # no shuffle
model = create_model(opt)
model.setup(opt)
model.eval()
print('Loading model %s' % opt.model)
testdata = ['manga_paper']
# fake_sty = model.get_z_random(1, 64, truncation=True)
opt.dataset_mode = 'singleSr'
for folder in testdata:
opt.folder = folder
# create dataset
dataset = create_dataset(opt)
web_dir = os.path.join(opt.results_dir, opt.folder + '_Sr2Co')
webpage = html.HTML(web_dir, 'Training = %s, Phase = %s, Class =%s' % (opt.name, opt.phase, opt.name))
# fake_sty = model.get_z_random(1, 64, truncation=True)
for i, data in enumerate(islice(dataset, opt.num_test)):
h = data['h']
w = data['w']
model.set_input(data)
fake_sty = model.get_z_random(1, 64, truncation=True, tvalue=1.25)
fake_B, SCR, line = model.forward(AtoB=False, sty=fake_sty)
images=[fake_B[:,:,:h,:w]]
names=['color']
img_path = 'input_%3.3d' % i
save_images(webpage, images, names, img_path, aspect_ratio=opt.aspect_ratio, width=opt.crop_size)
webpage.save()
testdata = ['western_paper']
opt.dataset_mode = 'singleCo'
for folder in testdata:
opt.folder = folder
# create dataset
dataset = create_dataset(opt)
web_dir = os.path.join(opt.results_dir, opt.folder + '_Sr2Co')
webpage = html.HTML(web_dir, 'Training = %s, Phase = %s, Class =%s' % (opt.name, opt.phase, opt.name))
for i, data in enumerate(islice(dataset, opt.num_test)):
h = data['h']
w = data['w']
model.set_input(data)
fake_B, fake_B2, SCR = model.forward(AtoB=True)
images=[fake_B2[:,:,:h,:w]]
names=['manga']
img_path = 'input_%3.3d' % i
save_images(webpage, images, names, img_path, aspect_ratio=opt.aspect_ratio, width=opt.crop_size)
webpage.save()
|