149 lines
4.8 KiB
Python
149 lines
4.8 KiB
Python
import base64
|
|
import copy
|
|
import logging
|
|
import sys
|
|
import time
|
|
from contextlib import contextmanager
|
|
from typing import Dict, List, Optional, Sequence, Tuple, Type, Union
|
|
|
|
import numpy as np
|
|
|
|
# TypedDict was added to typing in 3.8. Use typing_extensions for <3.8
|
|
if sys.version_info >= (3, 8):
|
|
from typing import TypedDict # pylint: disable=no-name-in-module
|
|
else:
|
|
from typing_extensions import TypedDict
|
|
|
|
|
|
class Timer(object):
|
|
def __init__(self, name: str):
|
|
self.name: str = name
|
|
self.start: Optional[float] = None
|
|
self.end: Optional[float] = None
|
|
|
|
def __enter__(self):
|
|
self.start = time.time()
|
|
|
|
def __exit__(self, type_, value, traceback):
|
|
self.end = time.time()
|
|
logging.debug("%s time: %s", self.name, self.end - self.start)
|
|
|
|
|
|
JSONArrayType = TypedDict(
|
|
"JSONArrayType",
|
|
{"@type": str, "base64": str, "dtype": str, "shape": Tuple[int, ...]},
|
|
)
|
|
|
|
|
|
def deserialise_array_b64(
|
|
b64_string: str, dtype: Union[Type[np.dtype], str], shape: Tuple[int, ...]
|
|
):
|
|
flat_arr: np.ndarray = np.frombuffer(base64.b64decode(b64_string), dtype)
|
|
return flat_arr.reshape(shape)
|
|
|
|
|
|
def serialise_array_b64(npy_arr: np.ndarray) -> Tuple[str, str, Tuple[int, ...]]:
|
|
b64_string: str = base64.b64encode(npy_arr.tobytes()).decode("ascii")
|
|
dtype: str = str(npy_arr.dtype)
|
|
shape: Tuple[int, ...] = npy_arr.shape
|
|
return b64_string, dtype, shape
|
|
|
|
|
|
def ndarray_to_json(arr: np.ndarray) -> JSONArrayType:
|
|
if isinstance(arr, memoryview):
|
|
# We can transparently convert memoryview objects to arrays
|
|
# This comes in very handy for the lens shading table.
|
|
arr = np.array(arr)
|
|
b64_string, dtype, shape = serialise_array_b64(arr)
|
|
return {"@type": "ndarray", "dtype": dtype, "shape": shape, "base64": b64_string}
|
|
|
|
|
|
def json_to_ndarray(json_dict: JSONArrayType):
|
|
if not json_dict.get("@type") != "ndarray":
|
|
logging.warning("No valid @type attribute found. Conversion may fail.")
|
|
for required_param in ("dtype", "shape", "base64"):
|
|
if not json_dict.get(required_param):
|
|
raise KeyError(f"Missing required key {required_param}")
|
|
|
|
b64_string: Optional[str] = json_dict.get("base64")
|
|
dtype: Optional[str] = json_dict.get("dtype")
|
|
shape: Optional[Tuple[int, ...]] = json_dict.get("shape")
|
|
|
|
if b64_string and dtype and shape:
|
|
return deserialise_array_b64(b64_string, dtype, shape)
|
|
else:
|
|
raise ValueError("Required parameters for decoding are missing")
|
|
|
|
|
|
@contextmanager
|
|
def set_properties(obj, **kwargs):
|
|
"""A context manager to set, then reset, certain properties of an object.
|
|
|
|
The first argument is the object, subsequent keyword arguments are properties
|
|
of said object, which are set initially, then reset to their previous values.
|
|
"""
|
|
saved_properties = {}
|
|
for k in kwargs.keys():
|
|
try:
|
|
saved_properties[k] = getattr(obj, k)
|
|
except AttributeError:
|
|
print(
|
|
"Warning: could not get {} on {}. This property will not be restored!".format(
|
|
k, obj
|
|
)
|
|
)
|
|
for k, v in kwargs.items():
|
|
setattr(obj, k, v)
|
|
try:
|
|
yield
|
|
finally:
|
|
for k, v in saved_properties.items():
|
|
setattr(obj, k, v)
|
|
|
|
|
|
def axes_to_array(
|
|
coordinate_dictionary: Dict[str, Optional[int]],
|
|
axis_keys: Sequence[str] = ("x", "y", "z"),
|
|
base_array: Optional[List[int]] = None,
|
|
asint: bool = True,
|
|
) -> List[int]:
|
|
"""Takes key-value pairs of a JSON value, and maps onto an array
|
|
|
|
This is designed to take a dictionary like `{"x": 1, "y":2, "z":3}`
|
|
and return a list like `[1, 2, 3]` to convert between the argument
|
|
format expected by most of our stages, and the usual argument
|
|
format in JSON.
|
|
|
|
`axis_keys` is an ordered sequence of key names to extract from
|
|
the input dictionary.
|
|
|
|
`base_array` specifies a default value for each axis. It must
|
|
have the same length as `axis_keys`.
|
|
|
|
`asint` casts values to integers if it is `True` (default).
|
|
|
|
Missing keys, or keys that have a `None` value will be left
|
|
at the specified default value, or zero if none is specified.
|
|
"""
|
|
# If no base array is given
|
|
if not base_array:
|
|
# Create an array of zeros
|
|
base_array = [0] * len(axis_keys)
|
|
else:
|
|
# Create a copy of the passed base_array
|
|
base_array = copy.copy(base_array)
|
|
|
|
# Do the mapping
|
|
for axis, key in enumerate(axis_keys):
|
|
if key in coordinate_dictionary:
|
|
value = coordinate_dictionary[key]
|
|
if value is None:
|
|
# Values set to None should be treated as if they
|
|
# are missing
|
|
# i.e. we leave the default value in place.
|
|
break
|
|
if asint:
|
|
value = int(value)
|
|
base_array[axis] = value
|
|
|
|
return base_array
|