39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
import os
|
|
import torch
|
|
from .saver import Saver
|
|
from .options import Options
|
|
from .model import CSDMT
|
|
from .dataset import MakeupDataset
|
|
import warnings
|
|
warnings.filterwarnings("ignore")
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
|
|
|
|
|
|
def eval_pair():
|
|
# parse options
|
|
parser = Options()
|
|
opts = parser.parse()
|
|
opts.phase = 'test_pair'
|
|
opts.data_root='./examples/images'
|
|
opts.img_dir='./inference_results'
|
|
# data loader
|
|
print('\n--- load dataset ---')
|
|
dataset = MakeupDataset(opts)
|
|
train_loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False,num_workers=0)
|
|
|
|
# model
|
|
print('\n--- load model ---')
|
|
model = CSDMT(opts)
|
|
ep0, total_it = model.resume(os.path.join(opts.model_dir, 'makeup_00499.pth'))
|
|
model.eval()
|
|
print('start test pair')
|
|
# saver for display and output
|
|
saver = Saver(opts)
|
|
for iter, data in enumerate(train_loader):
|
|
with torch.no_grad():
|
|
saver.write_test_pair_img(iter, model,data)
|
|
|
|
if __name__=='__main__':
|
|
eval_pair()
|
|
|