From c1117852d87929701288fdd2ee5a12a3ed94eb13 Mon Sep 17 00:00:00 2001
From: iperov <lepersorium@gmail.com>
Date: Sat, 30 Mar 2019 21:52:36 +0400
Subject: [PATCH] +imagelib.normalize_channels

---
 imagelib/__init__.py |  4 +++-
 imagelib/common.py   | 19 +++++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)
 create mode 100644 imagelib/common.py

diff --git a/imagelib/__init__.py b/imagelib/__init__.py
index 64e5a90..ab2ed81 100644
--- a/imagelib/__init__.py
+++ b/imagelib/__init__.py
@@ -18,4 +18,6 @@ from .color_transfer import color_hist_match
 from .color_transfer import reinhard_color_transfer
 from .color_transfer import linear_color_transfer
 
-from .DCSCN import DCSCN
\ No newline at end of file
+from .DCSCN import DCSCN
+
+from .common import normalize_channels
\ No newline at end of file
diff --git a/imagelib/common.py b/imagelib/common.py
new file mode 100644
index 0000000..fc723f1
--- /dev/null
+++ b/imagelib/common.py
@@ -0,0 +1,19 @@
+def normalize_channels(img, target_channels):
+    img_shape_len = len(img.shape)
+    if img_shape_len == 2:
+        h, w = img.shape
+        c = 0
+    elif img_shape_len == 3:
+        h, w, c = img.shape
+    else:
+        raise ValueError("normalize: incorrect image dimensions.")
+        
+    if c == 0 and target_channels > 0:
+        img = img[...,np.newaxis]        
+    if c == 1 and target_channels > 1:
+        img = np.repeat (img, target_channels, -1)   
+    if c > target_channels:         
+        img = img[...,0:target_channels]
+        c = target_channels
+        
+    return img
\ No newline at end of file