瀏覽代碼

add is_gray_image

quarrying 4 年之前
父節點
當前提交
ffb403d6e3
共有 1 個文件被更改,包括 23 次插入0 次删除
  1. 23 0
      khandy/image/misc.py

+ 23 - 0
khandy/image/misc.py

@@ -62,3 +62,26 @@ def stack_image_list(image_list, dtype=np.float32):
     return blob
     return blob
     
     
     
     
+def is_gray_image(image, tol=3):
+    if image.ndim == 2:
+        return True
+    elif image.ndim == 3:
+        num_channels = image.shape[-1]
+        if num_channels == 1:
+            return True
+        elif num_channels == 4:
+            rgb = cv2.cvtColor(image, cv2.COLOR_BGRA2BGR)
+            gray = cv2.cvtColor(rgb, cv2.COLOR_BGR2GRAY)
+            gray3 = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
+            mae = np.mean(cv2.absdiff(rgb, gray3))
+            return mae <= tol
+        elif num_channels == 3:
+            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+            gray3 = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
+            mae = np.mean(cv2.absdiff(image, gray3))
+            return mae <= tol
+        else:
+            return False
+    else:
+        return False
+