1
1
import collections
2
+ import math
2
3
from enum import IntEnum
3
4
4
5
import cv2
7
8
from core import imagelib
8
9
from facelib import FaceType , LandmarksProcessor
9
10
11
+
10
12
class SampleProcessor (object ):
11
13
class SampleType (IntEnum ):
12
14
NONE = 0
@@ -114,8 +116,8 @@ def get_eyes_mask():
114
116
if sample_type == SPST .FACE_IMAGE or sample_type == SPST .FACE_MASK :
115
117
if not is_face_sample :
116
118
raise ValueError ("face_samples should be provided for sample_type FACE_*" )
117
-
118
- if is_face_sample :
119
+
120
+ if sample_type == SPST . FACE_IMAGE or sample_type == SPST . FACE_MASK :
119
121
face_type = opts .get ('face_type' , None )
120
122
face_mask_type = opts .get ('face_mask_type' , SPFMT .NONE )
121
123
@@ -125,7 +127,6 @@ def get_eyes_mask():
125
127
if face_type > sample .face_type :
126
128
raise Exception ('sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.' % (sample .filename , sample .face_type , face_type ) )
127
129
128
- if sample_type == SPST .FACE_IMAGE or sample_type == SPST .FACE_MASK :
129
130
130
131
if sample_type == SPST .FACE_MASK :
131
132
@@ -156,7 +157,7 @@ def get_eyes_mask():
156
157
img = cv2 .resize ( img , (resolution , resolution ), cv2 .INTER_CUBIC )
157
158
158
159
img = imagelib .warp_by_params (params_per_resolution [resolution ], img , warp , transform , can_flip = True , border_replicate = border_replicate , cv2_inter = cv2 .INTER_LINEAR )
159
-
160
+
160
161
if len (img .shape ) == 2 :
161
162
img = img [...,None ]
162
163
@@ -175,11 +176,11 @@ def get_eyes_mask():
175
176
else :
176
177
if w != resolution :
177
178
img = cv2 .resize ( img , (resolution , resolution ), cv2 .INTER_CUBIC )
178
- img = imagelib .warp_by_params (params_per_resolution [resolution ], img , warp , transform , can_flip = True , border_replicate = border_replicate )
179
179
180
+ img = imagelib .warp_by_params (params_per_resolution [resolution ], img , warp , transform , can_flip = True , border_replicate = border_replicate )
181
+
180
182
img = np .clip (img .astype (np .float32 ), 0 , 1 )
181
-
182
-
183
+
183
184
184
185
# Apply random color transfer
185
186
if ct_mode is not None and ct_sample is not None :
@@ -273,17 +274,16 @@ def get_eyes_mask():
273
274
l = np .clip (l , 0.0 , 1.0 )
274
275
out_sample = l
275
276
elif sample_type == SPST .PITCH_YAW_ROLL or sample_type == SPST .PITCH_YAW_ROLL_SIGMOID :
276
- pitch_yaw_roll = sample .get_pitch_yaw_roll ()
277
-
278
- if params ['flip' ]:
277
+ pitch ,yaw ,roll = sample .get_pitch_yaw_roll ()
278
+ if params_per_resolution [resolution ]['flip' ]:
279
279
yaw = - yaw
280
280
281
281
if sample_type == SPST .PITCH_YAW_ROLL_SIGMOID :
282
282
pitch = np .clip ( (pitch / math .pi ) / 2.0 + 0.5 , 0 , 1 )
283
283
yaw = np .clip ( (yaw / math .pi ) / 2.0 + 0.5 , 0 , 1 )
284
284
roll = np .clip ( (roll / math .pi ) / 2.0 + 0.5 , 0 , 1 )
285
285
286
- out_sample = (pitch , yaw , roll )
286
+ out_sample = (pitch , yaw )
287
287
else :
288
288
raise ValueError ('expected sample_type' )
289
289
0 commit comments