Browse Source

refactor crop_or_pad and update somewhere

quarrying 4 years ago
parent
commit
b0fce33c48
2 changed files with 23 additions and 39 deletions
  1. 0 11
      khandy/image/align_and_crop.py
  2. 23 28
      khandy/image/crop_or_pad.py

+ 0 - 11
khandy/image/align_and_crop.py

@@ -1,8 +1,6 @@
 import cv2
 import numpy as np
 
-from .crop_or_pad import crop_or_pad as _crop_or_pad
-
 
 def get_similarity_transform(src_pts, dst_pts):
     """Get similarity transform matrix from src_pts to dst_pts
@@ -40,7 +38,6 @@ def get_similarity_transform(src_pts, dst_pts):
     
     
 def align_and_crop(image, landmarks, std_landmarks, align_size, 
-                   crop_size=None, crop_center=None,
                    return_transform_matrix=False):
     landmarks = np.asarray(landmarks)
     std_landmarks = np.asarray(std_landmarks)
@@ -49,14 +46,6 @@ def align_and_crop(image, landmarks, std_landmarks, align_size,
     landmarks_ex = np.pad(landmarks, ((0,0),(0,1)), mode='constant', constant_values=1)
     dst_landmarks = np.dot(landmarks_ex, xform_matrix[:2,:].T)
     dst_image = cv2.warpAffine(image, xform_matrix[:2,:], dsize=align_size)
-    if crop_size is not None:
-        crop_center_ex = (crop_center[0], crop_center[1], 1)
-        aligned_crop_center = np.dot(xform_matrix, crop_center_ex)
-        dst_image = _crop_or_pad(dst_image, crop_size, aligned_crop_center)
-        
-        crop_begin_x = int(round(aligned_crop_center[0] - crop_size[0] / 2.0))
-        crop_begin_y = int(round(aligned_crop_center[1] - crop_size[1] / 2.0))
-        dst_landmarks -= np.asarray([[crop_begin_x, crop_begin_y]])
     if return_transform_matrix:
         return dst_image, dst_landmarks, xform_matrix
     else:

+ 23 - 28
khandy/image/crop_or_pad.py

@@ -2,51 +2,47 @@ import cv2
 import numpy as np
 
 
-def crop_or_pad(image, crop_size, crop_center=None, pad_val=None):
+def crop_or_pad(image, x_min, y_min, x_max, y_max, pad_val=None):
     """
     References:
         tf.image.resize_image_with_crop_or_pad
     """
     assert image.ndim in [2, 3]
+    assert isinstance(x_min, int) and isinstance(y_min, int)
+    assert isinstance(x_max, int) and isinstance(y_max, int)
+    assert (x_min <= x_max) and (y_min <= y_max)
     
     src_height, src_width = image.shape[:2]
+    dst_height, dst_width = y_max - y_min + 1, x_max - x_min + 1
     channels = 1 if image.ndim == 2 else image.shape[2]
-    dst_height, dst_width = crop_size[1], crop_size[0]
-    if crop_center is None:
-        crop_center = [src_width // 2, src_height // 2]
+    
     if pad_val is not None:
         if isinstance(pad_val, (int, float)):
             pad_val = [pad_val for _ in range(channels)]
         assert len(pad_val) == channels
         
-    crop_begin_x = int(round(crop_center[0] - dst_width / 2.0))
-    crop_begin_y = int(round(crop_center[1] - dst_height / 2.0))
-    
-    src_begin_x = max(0, crop_begin_x)
-    src_begin_y = max(0, crop_begin_y)
-    src_end_x = min(src_width, crop_begin_x + dst_width)
-    src_end_y = min(src_height, crop_begin_y + dst_height)
-    dst_begin_x = max(0, -crop_begin_x)
-    dst_begin_y = max(0, -crop_begin_y)
-    dst_end_x = dst_begin_x + src_end_x - src_begin_x
-    dst_end_y = dst_begin_y + src_end_y - src_begin_y
+    src_x_begin = max(x_min, 0)
+    src_y_begin = max(y_min, 0, )
+    src_x_end = min(x_max + 1, src_width)
+    src_y_end = min(y_max + 1, src_height)
+    dst_x_begin = src_x_begin - x_min
+    dst_y_begin = src_y_begin - y_min
+    dst_x_end = src_x_end - x_min
+    dst_y_end = src_y_end - y_min
     
     if image.ndim == 2: 
-        cropped_image_shape = (dst_height, dst_width)
+        dst_image_shape = (dst_height, dst_width)
     else:
-        cropped_image_shape = (dst_height, dst_width, channels)
+        dst_image_shape = (dst_height, dst_width, channels)
     if pad_val is None:
-        cropped = np.zeros(cropped_image_shape, image.dtype)
+        dst_image = np.zeros(dst_image_shape, image.dtype)
     else:
-        cropped = np.full(cropped_image_shape, pad_val, dtype=image.dtype)
-    if (src_end_x - src_begin_x <= 0) or (src_end_y - src_begin_y <= 0):
-        return cropped
-    else:
-        cropped[dst_begin_y: dst_end_y, dst_begin_x: dst_end_x, ...] = \
-            image[src_begin_y: src_end_y, src_begin_x: src_end_x, ...]
-        return cropped
-        
-        
+        dst_image = np.full(dst_image_shape, pad_val, dtype=image.dtype)
+    dst_image[dst_y_begin: dst_y_end, dst_x_begin: dst_x_end, ...] = \
+        image[src_y_begin: src_y_end, src_x_begin: src_x_end, ...]
+    return dst_image
+    
+    
 def crop_or_pad_coords(boxes, image_width, image_height):
     """
     References:
@@ -73,5 +69,4 @@ def crop_or_pad_coords(boxes, image_width, image_height):
                        dst_x_begin, dst_y_begin, dst_x_end, dst_y_end], axis=1)
     return coords
     
-