fixes

iperov · Mar 16, 2020 · f3b4658 · f3b4658
1 parent 2300da4
commit f3b4658
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 14 deletions.
diff --git a/core/imagelib/warp.py b/core/imagelib/warp.py
@@ -25,8 +25,8 @@ def gen_warp_params (w, flip, rotation_range=[-10,10], scale_range=[-0.5, 0.5],
 
     half_cell_size = cell_size // 2
 
-    mapx = cv2.resize(mapx, (w+cell_size,)*2 )[half_cell_size:-half_cell_size-1,half_cell_size:-half_cell_size-1].astype(np.float32)
-    mapy = cv2.resize(mapy, (w+cell_size,)*2 )[half_cell_size:-half_cell_size-1,half_cell_size:-half_cell_size-1].astype(np.float32)
+    mapx = cv2.resize(mapx, (w+cell_size,)*2 )[half_cell_size:-half_cell_size,half_cell_size:-half_cell_size].astype(np.float32)
+    mapy = cv2.resize(mapy, (w+cell_size,)*2 )[half_cell_size:-half_cell_size,half_cell_size:-half_cell_size].astype(np.float32)
 
     #random transform
     random_transform_mat = cv2.getRotationMatrix2D((w // 2, w // 2), rotation, scale)

diff --git a/samplelib/Sample.py b/samplelib/Sample.py
@@ -64,7 +64,7 @@ def __init__(self, sample_type=None,
 
     def get_pitch_yaw_roll(self):
         if self.pitch_yaw_roll is None:
-            self.pitch_yaw_roll = LandmarksProcessor.estimate_pitch_yaw_roll(landmarks, size=self.shape[1])
+            self.pitch_yaw_roll = LandmarksProcessor.estimate_pitch_yaw_roll(self.landmarks, size=self.shape[1])
         return self.pitch_yaw_roll
 
     def set_filename_offset_size(self, filename, offset, size):

diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py
@@ -1,4 +1,5 @@
 import collections
+import math
 from enum import IntEnum
 
 import cv2
@@ -7,6 +8,7 @@
 from core import imagelib
 from facelib import FaceType, LandmarksProcessor
 
+
 class SampleProcessor(object):
     class SampleType(IntEnum):
         NONE = 0
@@ -114,8 +116,8 @@ def get_eyes_mask():
                 if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK:
                     if not is_face_sample:    
                         raise ValueError("face_samples should be provided for sample_type FACE_*")
-                    
-                if is_face_sample:    
+
+                if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK:
                     face_type      = opts.get('face_type', None)
                     face_mask_type = opts.get('face_mask_type', SPFMT.NONE)
 
@@ -125,7 +127,6 @@ def get_eyes_mask():
                     if face_type > sample.face_type:
                         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) )
 
-                if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK:
 
                     if sample_type == SPST.FACE_MASK:                        
 
@@ -156,7 +157,7 @@ def get_eyes_mask():
                                     img = cv2.resize( img, (resolution, resolution), cv2.INTER_CUBIC )
 
                             img = imagelib.warp_by_params (params_per_resolution[resolution], img, warp, transform, can_flip=True, border_replicate=border_replicate, cv2_inter=cv2.INTER_LINEAR)
-                                
+
                         if len(img.shape) == 2:
                             img = img[...,None]
 
@@ -175,11 +176,11 @@ def get_eyes_mask():
                         else:
                             if w != resolution:
                                 img = cv2.resize( img, (resolution, resolution), cv2.INTER_CUBIC )
-                        img  = imagelib.warp_by_params (params_per_resolution[resolution], img,  warp, transform, can_flip=True, border_replicate=border_replicate)
 
+                        img  = imagelib.warp_by_params (params_per_resolution[resolution], img,  warp, transform, can_flip=True, border_replicate=border_replicate)
+
                         img = np.clip(img.astype(np.float32), 0, 1)
-
-
+
 
                         # Apply random color transfer                        
                         if ct_mode is not None and ct_sample is not None:
@@ -273,17 +274,16 @@ def get_eyes_mask():
                     l = np.clip(l, 0.0, 1.0)
                     out_sample = l
                 elif sample_type == SPST.PITCH_YAW_ROLL or sample_type == SPST.PITCH_YAW_ROLL_SIGMOID:
-                    pitch_yaw_roll = sample.get_pitch_yaw_roll()
-
-                    if params['flip']:
+                    pitch,yaw,roll = sample.get_pitch_yaw_roll()
+                    if params_per_resolution[resolution]['flip']:
                         yaw = -yaw
 
                     if sample_type == SPST.PITCH_YAW_ROLL_SIGMOID:
                         pitch = np.clip( (pitch / math.pi) / 2.0 + 0.5, 0, 1)
                         yaw   = np.clip( (yaw / math.pi) / 2.0 + 0.5, 0, 1)
                         roll  = np.clip( (roll / math.pi) / 2.0 + 0.5, 0, 1)
 
-                    out_sample = (pitch, yaw, roll)
+                    out_sample = (pitch, yaw)
                 else:
                     raise ValueError ('expected sample_type')