https://github.com/IGLICT/DeepFaceDrawing-Jittor.git
Tip revision: 6319787fe99a47833b04fee43b98511ad4e98077 authored by Shuyu Chen on 06 March 2021, 06:53:00 UTC
Update install.sh
Update install.sh
Tip revision: 6319787
test_model.py
import time
from CombineModel_jt import CombineModel
import cv2
import numpy as np
import os
import glob
from tqdm import tqdm
import jittor as jt
jt.flags.use_cuda = 1
#models for face/eye1/eye2/nose/mouth
combine_model = CombineModel()
print('start')
fileRoot = './test'
images_path = sorted(glob.glob(fileRoot+r"/*"))
params = [
[0.80, 0.63, 1.0, 0.88, 0.93, 1],
[1.0, 1.0, 1.0, 1.0, 0.84, 0],
[0.1, 0.39, 0.58, 0.63, 0.49, 1],
[1.0, 1.0, 1.0, 1.0, 1.0, 1],
[0.78, 1.0, 1.0, 1.0, 0.79, 1]
]
i = 0
for x,fileName in enumerate(images_path):
print('Input file:',fileName)
mat_img = cv2.imread(fileName)
mat_img = cv2.resize(mat_img, (512, 512), interpolation=cv2.INTER_CUBIC)
mat_img = cv2.cvtColor(mat_img, cv2.COLOR_RGB2BGR)
sketch = (mat_img).astype(np.uint8)
combine_model.sex = params[i][5]
#666
combine_model.part_weight['eye1'] = params[i][0]
combine_model.part_weight['eye2'] = params[i][1]
combine_model.part_weight['nose'] = params[i][2]
combine_model.part_weight['mouth'] = params[i][3]
combine_model.part_weight[''] = params[i][4]
combine_model.predict_shadow(mat_img)
output_file = 'ori'+ str(x) +'.jpg'
print('Output file:',output_file)
cv2.imwrite(output_file,cv2.cvtColor(combine_model.generated, cv2.COLOR_BGR2RGB))
i = i + 1
jt.gc()