video.normalize: add dtype argument, support complex numbers

Author: sitic

docs/tutorials/io.ipynb

@@ -282,20 +282,12 @@
    },
    "outputs": [],
    "source": [
-    "om.save_video('video.npy', video)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "tags": [
-     "skip-execution"
-    ]
-   },
-   "outputs": [],
-   "source": [
-    "om.save_video('video.mat', video)"
+    "# NumPy NPY format\n",
+    "om.save_video('video.npy', video)\n",
+    "# MATLAB\n",
+    "om.save_video('video.mat', video)\n",
+    "# TIFF image stack\n",
+    "om.save_video('video.tif', video)"
    ]
   },
   {
@@ -350,13 +342,8 @@
    },
    "outputs": [],
    "source": [
-    "import numpy as np\n",
-    "\n",
-    "om.print_properties(video)\n",
-    "video = video.astype(np.float32)\n",
-    "video = (video-np.min(video))/np.max(video)*255 # alternatively *65536 for 16-bit\n",
-    "om.print_properties(video)\n",
-    "video = video.astype(np.uint8) # alternatively np.uint16 for 16-bit\n",
+    "# Convert to uint8 with range [0-255]\n",
+    "video = om.video.normalize(video, ymin=0, ymax=255, dtype=\"uint8\")\n",
     "om.print_properties(video)"
    ]
   },
@@ -431,7 +418,7 @@
    "outputs": [],
    "source": [
     "video_waves = om.video.normalize_pixelwise(video)\n",
-    "om.export_video_with_overlay(\"video.mp4\", video, video_waves)"
+    "om.video.export_video_with_overlay(\"video.mp4\", video, video_waves)"
    ]
   },
   {

optimap/video/_filters.py

@@ -7,14 +7,15 @@
 from ..utils import _print
 
 
-def normalize(video: np.ndarray, ymin=0, ymax=1, vmin=None, vmax=None, clip=True):
-    """Normalizes video to interval [``ymin``, ``ymax``]. If ``vmin`` or ``vmax`` are specified,
-    the normalization is done using these values and the resulting video will be clipped.
+def normalize(array: np.ndarray, ymin=0, ymax=1, vmin=None, vmax=None, clip=True, dtype=np.float32):
+    """Normalizes array (video, image, ...) to interval [``ymin``, ``ymax``].
+    
+    If ``vmin`` or ``vmax`` are specified, the normalization is done using these values and the resulting array will be clipped.
 
     Parameters
     ----------
-    video : {t, x, y} ndarray
-        The input video to be normalized.
+    array : ndarray
+        The input array to be normalized.
     ymin : float, optional
         Minimum value of the resulting video, by default 0
     ymax : float, optional
@@ -28,26 +29,37 @@ def normalize(video: np.ndarray, ymin=0, ymax=1, vmin=None, vmax=None, clip=True
     clip : bool, optional
         If True, the resulting video will be clipped to [``ymin``, ``ymax``], by default True
         Only applies if ``vmin`` or ``vmax`` are specified.
+    dtype : type, optional
+        Data type of the resulting array, by default np.float32
 
     Returns
     -------
-    {t, x, y} ndarray
-        Normalized video.
+    ndarray
+        Normalized array/video/image.
     """
     _print(f"normalizing video to interval [{ymin}, {ymax}] ...")
-    video = video.astype("float32")
+    dtype = np.dtype(dtype)
     do_clip = clip and (vmin is not None or vmax is not None)
+
+    if not (np.issubdtype(array.dtype, np.floating)
+            or np.issubdtype(array.dtype, np.complexfloating)):
+        array = array.astype(np.float32)
+
     if vmin is None:
-        vmin = np.nanmin(video)
+        vmin = np.nanmin(array)
     if vmax is None:
-        vmax = np.nanmax(video)
+        vmax = np.nanmax(array)
 
-    eps = np.finfo(np.float32).eps
-    video = (video - vmin) / (vmax - vmin + eps) * (ymax - ymin) + ymin
+    eps = np.finfo(array.dtype).eps
+    array = (array - vmin) / (vmax - vmin + eps) * (ymax - ymin) + ymin
 
     if do_clip:
-        video = np.clip(video, ymin, ymax)
-    return video
+        array = np.clip(array, ymin, ymax)
+    
+    if dtype == array.dtype:
+        return array
+    else:
+        return array.astype(dtype)
 
 
 def normalize_pixelwise(video: np.ndarray, ymin=0, ymax=1):

tests/test_video.py

@@ -12,6 +12,13 @@ def test_normalize():
     assert np.all(out >= -0.5)
     assert np.all(out <= 0.5)
 
+    vid = vid.astype(np.float64)
+    out = om.video.normalize(vid, dtype="float64")
+    assert out.dtype == np.float64
+
+    out = om.video.normalize(vid, ymax=255, dtype="uint8")
+    assert out.dtype == np.uint8
+
 def test_normalize_pixelwise():
     vid = (np.random.random((10, 128, 128)) * 8000).astype(np.uint16)
     out = om.video.normalize_pixelwise(vid, -0.5, 0.5)