Source code for supervisely.geometry.mask_3d

# coding: utf-8

# docs
from __future__ import annotations

import base64
import gzip
import tempfile
from typing import Dict, List, Literal, Optional, Tuple, Union

import nrrd
import numpy as np

from supervisely import logger
from supervisely._utils import unwrap_if_numpy
from supervisely.geometry.constants import (
    CLASS_ID,
    CREATED_AT,
    DATA,
    GEOMETRY_SHAPE,
    GEOMETRY_TYPE,
    ID,
    LABELER_LOGIN,
    MASK_3D,
    SPACE_ORIGIN,
    UPDATED_AT,
)
from supervisely.geometry.geometry import Geometry
from supervisely.io.fs import get_file_ext, get_file_name, remove_dir
from supervisely.io.json import JsonSerializable

if not hasattr(np, "bool"):
    np.bool = np.bool_


class PointVolume(JsonSerializable):
    """
    PointVolume (x, y, z) determines position of Mask3D. It locates the first sample.
    :class:`PointVolume<PointVolume>` object is immutable.

    :param x: Position of PointVolume object on X-axis.
    :type x: int or float
    :param y: Position of PointVolume object on Y-axis.
    :type y: int or float
    :param z: Position of PointVolume object on Z-axis.
    :type z: int or float
    :Usage example:

     .. code-block:: python

        import supervisely as sly

        x = 100
        y = 200
        z = 2
        loc = sly.PointVolume(x, y, z)
    """

    def __init__(self, x: Union[int, float], y: Union[int, float], z: Union[int, float]):
        self._x = round(unwrap_if_numpy(x))
        self._y = round(unwrap_if_numpy(y))
        self._z = round(unwrap_if_numpy(z))

    @property
    def x(self) -> int:
        """
        Position of PointVolume on X-axis.

        :return: X of PointVolume
        :rtype: :class:`int`
        :Usage example:

         .. code-block:: python

            print(loc.x)
            # Output: 100
        """
        return self._x

    @property
    def y(self) -> int:
        """
        Position of PointVolume on Y-axis.

        :return: Y of PointVolume
        :rtype: :class:`int`

        :Usage example:

         .. code-block:: python

            print(loc.y)
            # Output: 200
        """
        return self._y

    @property
    def z(self) -> int:
        """
        Position of PointVolume on Z-axis.

        :return: Z of PointVolume
        :rtype: :class:`int`
        :Usage example:

         .. code-block:: python

            print(loc.z)
            # Output: 2
        """
        return self._z

    def to_json(self) -> Dict:
        """
        Convert the PointVolume to a json dict.

        :return: Json format as a dict
        :rtype: :class:`dict`
        :Usage example:

         .. code-block:: python

            loc_json = loc.to_json()
            print(loc_json)
            # Output: {
            #           "space_origin": [
            #                            200,
            #                            200,
            #                            100
            #                           ]
            #         }

        """

        packed_obj = {SPACE_ORIGIN: [self.x, self.y, self.z]}
        return packed_obj

    @classmethod
    def from_json(cls, packed_obj) -> PointVolume:
        """
        Convert a json dict to PointVolume.

        :param data: PointVolume in json format as a dict.
        :type data: dict
        :return: PointVolume object
        :rtype: :class:`PointVolume<PointVolume>`
        :Usage example:

         .. code-block:: python

            import supervisely as sly

            loc_json = {
                "space_origin": [
                                200,
                                200,
                                100,
                                ]
                        }

            loc = sly.PointVolume.from_json(loc_json)
        """
        return cls(
            x=packed_obj["space_origin"][0],
            y=packed_obj["space_origin"][1],
            z=packed_obj["space_origin"][2],
        )


[docs]class Mask3D(Geometry): """ Mask 3D geometry for a single :class:`Label<supervisely.annotation.label.Label>`. :class:`Mask3D<Mask3D>` object is immutable. :param data: Mask 3D mask data. Must be a numpy array with only 2 unique values: [0, 1] or [0, 255] or [False, True]. :type data: np.ndarray :param sly_id: Mask 3D ID in Supervisely server. :type sly_id: int, optional :param class_id: ID of :class:`ObjClass<supervisely.annotation.obj_class.ObjClass>` to which Mask 3D belongs. :type class_id: int, optional :param labeler_login: Login of the user who created Mask 3D. :type labeler_login: str, optional :param updated_at: Date and Time when Mask 3D was modified last. Date Format: Year:Month:Day:Hour:Minute:Seconds. Example: '2021-01-22T19:37:50.158Z'. :type updated_at: str, optional :param created_at: Date and Time when Mask 3D was created. Date Format is the same as in "updated_at" parameter. :type created_at: str, optional :raises: :class:`ValueError`, if data is not bool or no pixels set to True in data :Usage example: .. code-block:: python import supervisely as sly # Create simple Mask 3D mask3d = np.zeros((3, 3, 3), dtype=np.bool_) mask3d[0:2, 0:2, 0:2] = True shape = sly.Mask3D(mask3d) print(shape.data) # Output: # [[[ True True False] # [ True True False] # [False False False]] # [[ True True False] # [ True True False] # [False False False]] # [[False False False] # [False False False] # [False False False]]] """ def __init__( self, data: np.ndarray, sly_id: Optional[int] = None, class_id: Optional[int] = None, labeler_login: Optional[str] = None, updated_at: Optional[str] = None, created_at: Optional[str] = None, ): super().__init__( sly_id=sly_id, class_id=class_id, labeler_login=labeler_login, updated_at=updated_at, created_at=created_at, ) if not isinstance(data, np.ndarray): raise TypeError('Mask 3D "data" argument must be numpy array object!') data_dims = len(data.shape) if data_dims != 3: raise ValueError( f'Mask 3D "data" argument must be a 3-dimensional numpy array. Instead got {data_dims} dimensions' ) if data.dtype != np.bool: if list(np.unique(data)) not in [[0, 1], [0, 255]]: raise ValueError( f"Mask 3D mask data values must be one of: [0 1], [0 255], [False True]. Instead got {np.unique(data)}." ) if list(np.unique(data)) == [0, 1]: data = np.array(data, dtype=bool) elif list(np.unique(data)) == [0, 255]: data = np.array(data / 255, dtype=bool) self.data = data self._space_origin = None self._space = None self._space_directions = None
[docs] @staticmethod def geometry_name(): """Return geometry name""" return "mask_3d"
[docs] @staticmethod def from_file(figure, file_path: str): """ Load figure geometry from file. :param figure: Spatial figure :type figure: VolumeFigure :param file_path: Path to nrrd file with data :type file_path: str """ mask3d_data, mask3d_header = nrrd.read(file_path) figure.geometry.data = mask3d_data try: figure.geometry._space_origin = PointVolume( x=mask3d_header["space origin"][0], y=mask3d_header["space origin"][1], z=mask3d_header["space origin"][2], ) figure.geometry._space = mask3d_header["space"] figure.geometry._space_directions = mask3d_header["space directions"] except KeyError as e: header_keys = ["'space'", "'space directions'", "'space origin'"] if str(e) in header_keys: logger.warning( f"The Mask3D geometry for figure ID '{get_file_name(file_path)}' doesn't contain optional space attributes that have similar names to {', '.join(header_keys)}. To set the values for these attributes, you can use information from the Volume associated with this figure object." ) path_without_filename = "/".join(file_path.split("/")[:-1]) remove_dir(path_without_filename)
[docs] @classmethod def create_from_file(cls, file_path: str) -> Mask3D: """ Creates Mask3D geometry from file. :param file_path: Path to nrrd file with data :type file_path: str """ mask3d_data, mask3d_header = nrrd.read(file_path) geometry = cls(data=mask3d_data) try: geometry._space_origin = PointVolume( x=mask3d_header["space origin"][0], y=mask3d_header["space origin"][1], z=mask3d_header["space origin"][2], ) geometry._space = mask3d_header["space"] geometry._space_directions = mask3d_header["space directions"] except KeyError as e: header_keys = ["'space'", "'space directions'", "'space origin'"] if str(e) in header_keys: logger.debug( f"The Mask3D geometry created from the file '{file_path}' doesn't contain optional space attributes that have similar names to {', '.join(header_keys)}. To set the values for these attributes, you can use information from the Volume associated with this figure object." ) return geometry
[docs] @classmethod def from_bytes(cls, geometry_bytes: bytes) -> Mask3D: """ Create a Mask3D geometry object from bytes. :param geometry_bytes: NRRD file represented as bytes. :type geometry_bytes: bytes :return: A Mask3D geometry object. :rtype: Mask3D """ with tempfile.NamedTemporaryFile(delete=True) as temp_file: temp_file.write(geometry_bytes) return cls.create_from_file(temp_file.name)
[docs] def to_json(self) -> Dict: """ Convert the Mask 3D to a json dict. :return: Json format as a dict :rtype: :class:`dict` :Usage example: .. code-block:: python import supervisely as sly mask = np.array([[[1 1 0] [1 1 0] [0 0 0]] [[1 1 0] [1 1 0] [0 0 0]] [[0 0 0] [0 0 0] [0 0 0]]], dtype=np.bool_) figure = sly.Mask3D(mask) figure_json = figure.to_json() print(json.dumps(figure_json, indent=4)) # Output: { # "mask_3d": { # "data": "eJzrDPBz5+WS4mJgYOD19HAJAtLMIMwIInOeqf8BUmwBPiGuQPr///9Lb86/C2QxlgT5BTM4PLuRBuTwebo4hlTMSa44sKHhISMDuxpTYrr03F6gDIOnq5/LOqeEJgDM5ht6", # }, # "shape": "mask_3d", # "geometryType": "mask_3d" # } """ res = { self._impl_json_class_name(): { DATA: self.data_2_base64(self.data), }, GEOMETRY_SHAPE: self.name(), GEOMETRY_TYPE: self.name(), } if self._space_origin: res[f"{self._impl_json_class_name()}"][f"{SPACE_ORIGIN}"] = [ self._space_origin.x, self._space_origin.y, self._space_origin.z, ] self._add_creation_info(res) return res
[docs] @classmethod def from_json(cls, json_data: Dict) -> Mask3D: """ Convert a json dict to Mask 3D. :param data: Mask in json format as a dict. :type data: dict :return: Mask3D object :rtype: :class:`Mask3D<Mask3D>` :Usage example: .. code-block:: python import supervisely as sly figure_json = { "mask_3d": { "data": "eJzrDPBz5+WS4mJgYOD19HAJAtLMIMwIInOeqf8BUmwBPiGuQPr///9Lb86/C2QxlgT5BTM4PLuRBuTwebo4hlTMSa44sKHhISMDuxpTYrr03F6gDIOnq5/LOqeEJgDM5ht6", }, "shape": "mask_3d", "geometryType": "mask_3d" } figure = sly.Mask3D.from_json(figure_json) """ if json_data == {}: return cls(data=np.zeros((3, 3, 3), dtype=np.bool_)) json_root_key = cls._impl_json_class_name() if json_root_key not in json_data: raise ValueError( "Data must contain {} field to create Mask3D object.".format(json_root_key) ) if DATA not in json_data[json_root_key]: raise ValueError( "{} field must contain {} field to create Mask3D object.".format( json_root_key, DATA ) ) data = cls.base64_2_data(json_data[json_root_key][DATA]) labeler_login = json_data.get(LABELER_LOGIN, None) updated_at = json_data.get(UPDATED_AT, None) created_at = json_data.get(CREATED_AT, None) sly_id = json_data.get(ID, None) class_id = json_data.get(CLASS_ID, None) instance = cls( data=data.astype(np.bool_), sly_id=sly_id, class_id=class_id, labeler_login=labeler_login, updated_at=updated_at, created_at=created_at, ) if SPACE_ORIGIN in json_data[json_root_key]: x, y, z = json_data[json_root_key][SPACE_ORIGIN] instance._space_origin = PointVolume(x=x, y=y, z=z) return instance
@classmethod def _impl_json_class_name(cls): """_impl_json_class_name""" return MASK_3D
[docs] @staticmethod def data_2_base64(data: np.ndarray) -> str: """ Convert numpy array to base64 encoded string. :param mask: Bool numpy array. :type mask: np.ndarray :return: Base64 encoded string :rtype: :class:`str` :Usage example: .. code-block:: python import supervisely as sly import os import nrrd address = 'https://app.supervisely.com/' token = 'Your Supervisely API Token' api = sly.Api(address, token) meta_json = api.project.get_meta(PROJECT_ID) meta = sly.ProjectMeta.from_json(meta_json) ann_json = api.volume.annotation.download_bulk(DATASET_ID, [VOLUME_ID]) figure_id = ann_json[0]["spatialFigures"][0]["id"] path_for_mesh = f"meshes/{figure_id}.nrrd" api.volume.figure.download_stl_meshes([figure_id], [path_for_mesh]) mask3d_data, _ = nrrd.read(path_for_mesh) encoded_string = sly.Mask3D.data_2_base64(mask3d_data) print(encoded_string) # 'H4sIAGWoWmQC/zPWMdYxrmFkZAAiIIAz4AAAE56ciyEAAAA=' """ shape_str = ",".join(str(dim) for dim in data.shape) data_str = data.tostring().decode("utf-8") combined_str = f"{shape_str}|{data_str}" compressed_string = gzip.compress(combined_str.encode("utf-8")) encoded_string = base64.b64encode(compressed_string).decode("utf-8") return encoded_string
[docs] @staticmethod def base64_2_data(encoded_string: str) -> np.ndarray: """ Convert base64 encoded string to numpy array. :param s: Input base64 encoded string. :type s: str :return: Bool numpy array :rtype: :class:`np.ndarray` :Usage example: .. code-block:: python import supervisely as sly encoded_string = 'H4sIAGWoWmQC/zPWMdYxrmFkZAAiIIAz4AAAE56ciyEAAAA=' figure_data = sly.Mask3D.base64_2_data(encoded_string) print(figure_data) # [[[1 1 0] # [1 1 0] # [0 0 0]] # [[1 1 0] # [1 1 0] # [0 0 0]] # [[0 0 0] # [0 0 0] # [0 0 0]]] """ compressed_bytes = base64.b64decode(encoded_string) decompressed_string = gzip.decompress(compressed_bytes).decode("utf-8") shape_str, data_str = decompressed_string.split("|") shape = tuple(int(dim) for dim in shape_str.split(",")) data_bytes = data_str.encode("utf-8") try: data = np.frombuffer(data_bytes, dtype=np.uint8).reshape(shape) except ValueError: logger.warn( "Can't reshape array with 'dtype=np.uint8'. Will try to automatically convert 'dtype=np.int16' to 'np.uint8' and reshape" ) data = np.frombuffer(data_bytes, dtype=np.int16) data = np.clip(data, 0, 1).astype(np.uint8) data = data.reshape(shape) logger.debug("Converted successfully!") return data
[docs] def add_mask_2d( self, mask_2d: np.ndarray, plane_name: Literal["axial", "sagittal", "coronal"], slice_index: int, origin: Optional[List[int]] = None, ): """ Draw a 2D mask on a 3D Mask. :param mask_2d: 2D array with a flat mask. :type mask_2d: np.ndarray :param plane_name: Name of the plane: "axial", "sagittal", "coronal". :type plane_name: str :param slice_index: Slice index of the volume figure. :type slice_index: int :param origin: (row, col) position. The top-left corner of the mask is located on the specified slice (optional). :type origin: Optional[List[int]], NoneType """ from supervisely.volume_annotation.plane import Plane Plane.validate_name(plane_name) mask_2d = np.fliplr(mask_2d) mask_2d = np.rot90(mask_2d, 1, (1, 0)) if plane_name == Plane.AXIAL: new_shape = self.data.shape[:2] elif plane_name == Plane.SAGITTAL: new_shape = self.data.shape[1:] elif plane_name == Plane.CORONAL: new_shape = self.data.shape[::2] if origin: x, y = origin # pylint: disable=possibly-used-before-assignment new_mask = np.zeros(new_shape, dtype=mask_2d.dtype) new_mask[x : x + mask_2d.shape[0], y : y + mask_2d.shape[1]] = mask_2d if plane_name == Plane.AXIAL: self.data[:, :, slice_index] = new_mask elif plane_name == Plane.SAGITTAL: self.data[slice_index, :, :] = new_mask elif plane_name == Plane.CORONAL: self.data[:, slice_index, :] = new_mask
@staticmethod def _bytes_from_nrrd(path: str) -> Tuple[str, bytes]: """ Read geometry from a file as bytes. The NRRD file must be named with a hexadecimal UUID value. Only NRRD files are supported. :param path: Path to the NRRD file containing geometry. :type path: str :return: A tuple containing the key hex value and geometry bytes, or (None, None) if the file is not found. :rtype: Tuple[str, bytes] """ if get_file_ext(path) == ".nrrd": key = get_file_name(path) with open(path, "rb") as file: geometry_bytes = file.read() return key, geometry_bytes else: return None, None @staticmethod def _bytes_from_nrrd_batch(paths: List[str]) -> Dict[str, bytes]: """ Read geometries from multiple files as bytes and map them to figure UUID hex values in a dictionary. The NRRD files must be named with a hexadecimal UUID value. Only NRRD files are supported. :param paths: Paths to the NRRD files containing geometry. :type paths: List[str] :return: A dictionary mapping figure UUID hex values to their respective geometries. :rtype: Dict[str, bytes] """ geometries_dict = {} for path in paths: key, geometry_bytes = Mask3D._bytes_from_nrrd(path) if key is None and geometry_bytes is None: continue geometries_dict[key] = geometry_bytes return geometries_dict