diff --git a/xlib/image/misc.py b/xlib/image/misc.py new file mode 100644 index 0000000..10f85ec --- /dev/null +++ b/xlib/image/misc.py @@ -0,0 +1,17 @@ +import numpy as np + +def get_NHWC_shape(img : np.ndarray): + """ + returns NHWC shape where missed dims are 1 + """ + ndim = img.ndim + if ndim not in [2,3,4]: + raise ValueError(f'img.ndim must be 2,3,4, not {ndim}.') + + if ndim == 2: + N, (H,W), C = 1, img.shape, 1 + elif ndim == 3: + N, (H,W,C) = 1, img.shape + else: + N,H,W,C = img.shape + return N,H,W,C \ No newline at end of file