瀏覽代碼

add strict argument for center_crop

quarrying 3 年之前
父節點
當前提交
05e7afded6
共有 1 個文件被更改,包括 18 次插入13 次删除
  1. 18 13
      khandy/image/crop_or_pad.py

+ 18 - 13
khandy/image/crop_or_pad.py

@@ -59,18 +59,6 @@ def crop_or_pad(image, x_min, y_min, x_max, y_max, border_value=0):
     return dst_image
     
     
-def center_crop(image, dst_width, dst_height):
-    assert image.ndim in [2, 3]
-    assert isinstance(dst_width, numbers.Integral) and isinstance(dst_height, numbers.Integral)
-    assert (image.shape[0] >= dst_height) and (image.shape[1] >= dst_width)
-
-    crop_top = (image.shape[0] - dst_height) // 2
-    crop_left = (image.shape[1] - dst_width) // 2
-    dst_image = image[crop_top: dst_height + crop_top, 
-                      crop_left: dst_width + crop_left, ...]
-    return dst_image
-    
-    
 def crop_or_pad_coords(boxes, image_width, image_height):
     """
     References:
@@ -100,4 +88,21 @@ def crop_or_pad_coords(boxes, image_width, image_height):
                        dst_heights, dst_widths], axis=0)
     return coords
     
-    
+    
+def center_crop(image, dst_width, dst_height, strict=True):
+    """
+    strict: 
+        when True, raise error if src size is less than dst size. 
+        when False, remain unchanged if src size is less than dst size, otherwise center crop.
+    """
+    assert image.ndim in [2, 3]
+    assert isinstance(dst_width, numbers.Integral) and isinstance(dst_height, numbers.Integral)
+    src_height, src_width = image.shape[:2]
+    if strict:
+        assert (src_height >= dst_height) and (src_width >= dst_width)
+
+    crop_top = max((src_height - dst_height) // 2, 0)
+    crop_left = max((src_width - dst_width) // 2, 0)
+    cropped = image[crop_top: dst_height + crop_top, 
+                    crop_left: dst_width + crop_left, ...]
+    return cropped