ome_zarr.scale: migrate scale.py to ome_zarr

Provides the `ome_zarr scale` CLI command as well as the
`ome_zarr.scale.Scaler` class for downsampling arrays as
multiscales.

originally: ome/omero-ms-zarr#61

.isort.cfg

@@ -1,2 +1,7 @@
 [settings]
-known_third_party = dask,numpy,pytest,requests,scipy,setuptools,skimage,vispy,zarr
+known_third_party = cv2,dask,numpy,pytest,requests,scipy,setuptools,skimage,vispy,zarr
+multi_line_output=6
+include_trailing_comma=False
+force_grid_wrap=0
+use_parentheses=True
+line_length=120

ome_zarr/cli.py

@@ -6,6 +6,7 @@
 from typing import List
 
 from .data import astronaut, coins, create_zarr
+from .scale import Scaler
 from .utils import download as zarr_download
 from .utils import info as zarr_info
 
@@ -31,11 +32,25 @@ def create(args: argparse.Namespace) -> None:
     config_logging(logging.INFO, args)
     if args.method == "coins":
         method = coins
+        label_name = "coins"
     elif args.method == "astronaut":
         method = astronaut
+        label_name = "circles"
     else:
         raise Exception(f"unknown method: {args.method}")
-    create_zarr(args.path, method=method)
+    create_zarr(args.path, method=method, label_name=label_name)
+
+
+def scale(args: argparse.Namespace) -> None:
+    scaler = Scaler(
+        copy_metadata=args.copy_metadata,
+        downscale=args.downscale,
+        in_place=args.in_place,
+        labeled=args.labeled,
+        max_layer=args.max_layer,
+        method=args.method,
+    )
+    scaler.scale(args.input_array, args.output_directory)
 
 
 def main(args: List[str] = None) -> None:
@@ -70,14 +85,39 @@ def main(args: List[str] = None) -> None:
     parser_download.add_argument("--name", default="")
     parser_download.set_defaults(func=download)
 
-    # coin
+    # create
     parser_create = subparsers.add_parser("create")
     parser_create.add_argument(
         "--method", choices=("coins", "astronaut"), default="coins"
     )
     parser_create.add_argument("path")
     parser_create.set_defaults(func=create)
 
+    parser_scale = subparsers.add_parser("scale")
+    parser_scale.add_argument("input_array")
+    parser_scale.add_argument("output_directory")
+    parser_scale.add_argument(
+        "--labeled",
+        action="store_true",
+        help="assert that the list of unique pixel values doesn't change",
+    )
+    parser_scale.add_argument(
+        "--copy-metadata",
+        action="store_true",
+        help="copies the array metadata to the new group",
+    )
+    parser_scale.add_argument(
+        "--method", choices=list(Scaler.methods()), default="nearest"
+    )
+    parser_scale.add_argument(
+        "--in-place", action="store_true", help="if true, don't write the base array"
+    )
+    parser_scale.add_argument("--downscale", type=int, default=2)
+    parser_scale.add_argument("--max_layer", type=int, default=4)
+    parser_scale.set_defaults(func=scale)
+
+    ns = parser.parse_args()
+
     if args is None:
         ns = parser.parse_args(sys.argv[1:])
     else:

ome_zarr/data.py

@@ -9,9 +9,9 @@
 from skimage.measure import label
 from skimage.morphology import closing, remove_small_objects, square
 from skimage.segmentation import clear_border
-from skimage.transform import pyramid_gaussian
 
 from .conversions import rgba_to_int
+from .scale import Scaler
 
 
 def coins() -> Tuple[List, List]:
@@ -35,32 +35,45 @@ def coins() -> Tuple[List, List]:
 
 
 def astronaut() -> Tuple[List, List]:
-    base = np.tile(data.astronaut(), (2, 2, 1))
-    gaussian = list(pyramid_gaussian(base, downscale=2, max_layer=4, multichannel=True))
-
-    pyramid = []
-    # convert each level of pyramid into 5D image (t, c, z, y, x)
-    for pixels in gaussian:
-        red = pixels[:, :, 0]
-        green = pixels[:, :, 1]
-        blue = pixels[:, :, 2]
-        # wrap to make 5D: (t, c, z, y, x)
-        pixels = np.array([np.array([red]), np.array([green]), np.array([blue])])
-        pixels = np.array([pixels])
-        pyramid.append(pixels)
-    return pyramid, []
+    scaler = Scaler()
+
+    pixels = rgb_to_5d(np.tile(data.astronaut(), (2, 2, 1)))
+    pyramid = scaler.nearest(pixels)
+
+    shape = list(pyramid[0].shape)
+    shape[1] = 1
+    label = np.zeros(shape)
+    make_circle(100, 100, 1, label[0, 0, 0, 200:300, 200:300])
+    make_circle(150, 150, 2, label[0, 0, 0, 250:400, 250:400])
+    labels = scaler.nearest(label)
+
+    return pyramid, labels
+
+
+def make_circle(h: int, w: int, value: int, target: np.ndarray) -> None:
+    x = np.arange(0, w)
+    y = np.arange(0, h)
+
+    cx = w // 2
+    cy = h // 2
+    r = min(w, h) // 2
+
+    mask = (x[np.newaxis, :] - cx) ** 2 + (y[:, np.newaxis] - cy) ** 2 < r ** 2
+    target[mask] = value
 
 
 def rgb_to_5d(pixels: np.ndarray) -> List:
     """convert an RGB image into 5D image (t, c, z, y, x)"""
     if len(pixels.shape) == 2:
-        channels = [[np.array(pixels)]]
+        stack = np.array([pixels])
+        channels = np.array([stack])
     elif len(pixels.shape) == 3:
-        size_c = pixels.shape(2)
-        channels = [np.array(pixels[:, :, c]) for c in range(size_c)]
+        size_c = pixels.shape[2]
+        channels = [np.array([pixels[:, :, c]]) for c in range(size_c)]
     else:
         assert False, f"expecting 2 or 3d: ({pixels.shape})"
-    return [np.array(channels)]
+    video = np.array([channels])
+    return video
 
 
 def write_multiscale(pyramid: List, group: zarr.Group) -> None:

ome_zarr/scale.py

@@ -0,0 +1,167 @@
+import inspect
+import logging
+import os
+from collections.abc import MutableMapping
+from dataclasses import dataclass
+from typing import Callable, Iterator, List
+
+import cv2
+import numpy as np
+import zarr
+from scipy.ndimage import zoom
+from skimage.transform import downscale_local_mean, pyramid_gaussian, pyramid_laplacian
+
+LOGGER = logging.getLogger("ome_zarr.scale")
+
+
+@dataclass
+class Scaler:
+
+    copy_metadata: bool = False
+    downscale: int = 2
+    in_place: bool = False
+    labeled: bool = False
+    max_layer: int = 4
+    method: str = "nearest"
+
+    @staticmethod
+    def methods() -> Iterator[str]:
+        funcs = inspect.getmembers(Scaler, predicate=inspect.isfunction)
+        for name, func in funcs:
+            if name in ("methods", "scale"):
+                continue
+            if name.startswith("_"):
+                continue
+            yield name
+
+    def scale(self, input_array: str, output_directory: str) -> None:
+
+        func = getattr(self, self.method, None)
+        if not func:
+            raise Exception
+
+        store = self._check_store(output_directory)
+        base = zarr.open_array(input_array)
+        pyramid = func(base)
+
+        if self.labeled:
+            self._assert_values(pyramid)
+
+        grp = self._create_group(store, base, pyramid)
+
+        if self.copy_metadata:
+            print(f"copying attribute keys: {list(base.attrs.keys())}")
+            grp.attrs.update(base.attrs)
+
+    def _check_store(self, output_directory: str) -> MutableMapping:
+        assert not os.path.exists(output_directory)
+        return zarr.DirectoryStore(output_directory)
+
+    def _assert_values(self, pyramid: List[np.ndarray]) -> None:
+        expected = set(np.unique(pyramid[0]))
+        print(f"level 0 {pyramid[0].shape} = {len(expected)} labels")
+        for i in range(1, len(pyramid)):
+            level = pyramid[i]
+            print(f"level {i}", pyramid[i].shape, len(expected))
+            found = set(np.unique(level))
+            if not expected.issuperset(found):
+                raise Exception(
+                    f"{len(found)} found values are not "
+                    "a subset of {len(expected)} values"
+                )
+
+    def _create_group(
+        self, store: MutableMapping, base: np.ndarray, pyramid: List[np.ndarray]
+    ) -> zarr.hierarchy.Group:
+        grp = zarr.group(store)
+        grp.create_dataset("base", data=base)
+        series = []
+        for i, dataset in enumerate(pyramid):
+            if i == 0:
+                path = "base"
+            else:
+                path = "%s" % i
+                grp.create_dataset(path, data=pyramid[i])
+            series.append({"path": path})
+        return grp
+
+    #
+    # Scaling methods
+    #
+
+    def nearest(self, base: np.ndarray) -> List[np.ndarray]:
+        def func(plane: np.ndarray, sizeY: int, sizeX: int) -> np.ndarray:
+            return cv2.resize(
+                plane,
+                dsize=(sizeY // self.downscale, sizeX // self.downscale),
+                interpolation=cv2.INTER_NEAREST,
+            )
+
+        return self._by_plane(base, func)
+
+    def gaussian(self, base: np.ndarray) -> List[np.ndarray]:
+        return list(
+            pyramid_gaussian(
+                base,
+                downscale=self.downscale,
+                max_layer=self.max_layer,
+                multichannel=False,
+            )
+        )
+
+    def laplacian(self, base: np.ndarray) -> List[np.ndarray]:
+        return list(
+            pyramid_laplacian(
+                base,
+                downscale=self.downscale,
+                max_layer=self.max_layer,
+                multichannel=False,
+            )
+        )
+
+    def local_mean(self, base: np.ndarray) -> List[np.ndarray]:
+        # FIXME: fix hard-coding
+        rv = [base]
+        for i in range(self.max_layer):
+            rv.append(
+                downscale_local_mean(
+                    rv[-1], factors=(1, 1, 1, self.downscale, self.downscale)
+                )
+            )
+        return rv
+
+    def zoom(self, base: np.ndarray) -> List[np.ndarray]:
+        rv = [base]
+        print(base.shape)
+        for i in range(self.max_layer):
+            print(i, self.downscale)
+            rv.append(zoom(base, self.downscale ** i))
+            print(rv[-1].shape)
+        return list(reversed(rv))
+
+    #
+    # Helpers
+    #
+
+    def _by_plane(
+        self, base: np.ndarray, func: Callable[[np.ndarray, int, int], np.ndarray],
+    ) -> np.ndarray:
+
+        assert 5 == len(base.shape)
+
+        rv = [base]
+        for i in range(self.max_layer):
+            fiveD = rv[-1]
+            # FIXME: fix hard-coding of dimensions
+            T, C, Z, Y, X = fiveD.shape
+
+            smaller = None
+            for t in range(T):
+                for c in range(C):
+                    for z in range(Z):
+                        out = func(fiveD[t][c][z][:], Y, X)
+                        if smaller is None:
+                            smaller = np.zeros((T, C, Z, out.shape[0], out.shape[1]))
+                        smaller[t][c][z] = out
+            rv.append(smaller)
+        return rv