-
Notifications
You must be signed in to change notification settings - Fork 140
/
test_model.py
49 lines (41 loc) · 1.33 KB
/
test_model.py
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
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()