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