samgis_core 1.2.1__py3-none-any.whl → 1.2.2__py3-none-any.whl

samgis_core/prediction_api/sam_onnx2.py

@@ -0,0 +1,178 @@
+"""
+Define a machine learning model executed by ONNX Runtime (https://onnxruntime.ai/)
+for Segment Anything (https://segment-anything.com).
+Modified from
+- https://github.com/vietanhdev/samexporter/
+- https://github.com/AndreyGermanov/sam_onnx_full_export/
+
+Copyright (c) 2023 Viet Anh Nguyen, Andrey Germanov
+Copyright (c) 2024-today Alessandro Trinca Tornidor
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+from numpy import array as np_array, pad as np_pad, zeros, ndarray, concatenate, float32
+from onnxruntime import get_available_providers, InferenceSession
+
+from samgis_core import app_logger
+from samgis_core.utilities.constants import DEFAULT_INPUT_SHAPE
+from samgis_core.utilities.type_hints import ListDict, EmbeddingPILImage, PIL_Image
+from samgis_core.utilities.utilities import convert_ndarray_to_pil, apply_coords
+
+
+class SegmentAnythingONNX2:
+    """
+    Segmentation model using SegmentAnything.
+    Compatible with onnxruntime 1.17.x and later
+    """
+
+    def __init__(self, encoder_model_path: str, decoder_model_path: str) -> None:
+        self.target_size = DEFAULT_INPUT_SHAPE[1]
+        self.input_size = DEFAULT_INPUT_SHAPE
+
+        # Load models
+        providers = get_available_providers()
+
+        # Pop TensorRT Runtime due to crashing issues
+        # TODO: Add back when TensorRT backend is stable
+        providers = [p for p in providers if p != "TensorrtExecutionProvider"]
+
+        if providers:
+            app_logger.info(
+                "Available providers for ONNXRuntime: %s", ", ".join(providers)
+            )
+        else:
+            app_logger.warning("No available providers for ONNXRuntime")
+        self.encoder_session = InferenceSession(
+            encoder_model_path, providers=providers
+        )
+        self.encoder_input_name = self.encoder_session.get_inputs()[0].name
+        app_logger.info("encoder_input_name:", self.encoder_input_name)
+        self.decoder_session = InferenceSession(
+            decoder_model_path, providers=providers
+        )
+
+    @staticmethod
+    def get_input_points(prompt: ListDict):
+        """Get input points"""
+        points = []
+        labels = []
+        for mark in prompt:
+            if mark["type"] == "point":
+                points.append(mark["data"])
+                labels.append(mark["label"])
+            elif mark["type"] == "rectangle":
+                points.append([mark["data"][0], mark["data"][1]])  # top left
+                points.append(
+                    [mark["data"][2], mark["data"][3]]
+                )  # bottom right
+                labels.append(2)
+                labels.append(3)
+        points, labels = np_array(points), np_array(labels)
+        return points, labels
+
+    def encode(self, img: PIL_Image | ndarray) -> EmbeddingPILImage:
+        """
+        Calculate embedding and metadata for a single image.
+
+        Args:
+            img: input image to embed
+
+        Returns:
+            embedding image dict useful to store and cache image embeddings
+        """
+        resized_image = self.preprocess_image(img)
+        padded_input_tensor = self.padding_tensor(resized_image)
+
+        # 2. GET IMAGE EMBEDDINGS USING IMAGE ENCODER
+        outputs = self.encoder_session.run(None, {"images": padded_input_tensor})
+        image_embedding = outputs[0]
+        img = convert_ndarray_to_pil(img)
+        return {
+            "image_embedding": image_embedding,
+            "original_size": img.size,
+            "resized_size": resized_image.size
+        }
+
+    def predict_masks(self, embedding: EmbeddingPILImage, prompt: ListDict):
+        """
+        Predict masks for a single image.
+        """
+        input_points, input_labels = self.get_input_points(prompt)
+
+        # Add a batch index, concatenate a padding point, and transform.
+        onnx_coord = concatenate([input_points, np_array([[0.0, 0.0]])], axis=0)[None, :, :]
+        onnx_label = concatenate([input_labels, np_array([-1])], axis=0)[None, :].astype(float32)
+
+        onnx_coord = apply_coords(onnx_coord, embedding)
+        orig_width, orig_height = embedding["original_size"]
+        app_logger.info(f"onnx_coord:{onnx_coord}.")
+
+        # RUN DECODER TO GET MASK
+        onnx_mask_input = zeros((1, 1, 256, 256), dtype=float32)
+        onnx_has_mask_input = zeros(1, dtype=float32)
+        output_masks, _, _ = self.decoder_session.run(None, {
+            "image_embeddings": embedding["image_embedding"],
+            "point_coords": onnx_coord,
+            "point_labels": onnx_label,
+            "mask_input": onnx_mask_input,
+            "has_mask_input": onnx_has_mask_input,
+            "orig_im_size": np_array([orig_height, orig_width], dtype=float32)
+        })
+        return output_masks
+
+    def preprocess_image(self, img: PIL_Image | ndarray):
+        """Resize image preserving aspect ratio using 'output_size_target' as a long side"""
+        from PIL import Image
+
+        app_logger.info(f"image type:{type(img)}, shape/size:{img.size}.")
+        try:
+            orig_width, orig_height = img.size
+        except TypeError:
+            img = Image.fromarray(img)
+            orig_width, orig_height = img.size
+
+        resized_height = self.target_size
+        resized_width = int(self.target_size / orig_height * orig_width)
+
+        if orig_width > orig_height:
+            resized_width = self.target_size
+            resized_height = int(self.target_size / orig_width * orig_height)
+
+        img = img.resize((resized_width, resized_height), Image.Resampling.BILINEAR)
+        return img
+
+    def padding_tensor(self, img: PIL_Image | ndarray):
+        # Prepare input tensor from image
+        tensor_input = np_array(img)
+        resized_width, resized_height = img.size
+
+        # Normalize input tensor numbers
+        mean = np_array([123.675, 116.28, 103.53])
+        std = np_array([[58.395, 57.12, 57.375]])
+        tensor_input = (tensor_input - mean) / std
+
+        # Transpose input tensor to shape (Batch,Channels,Height,Width
+        tensor_input = tensor_input.transpose(2, 0, 1)[None, :, :, :].astype(float32)
+
+        # Make image square self.target_size x self.target_size by padding short side by zeros
+        tensor_input = np_pad(tensor_input, ((0, 0), (0, 0), (0, 0), (0, self.target_size - resized_width)))
+        if resized_height < resized_width:
+            tensor_input = np_pad(tensor_input, ((0, 0), (0, 0), (0, self.target_size - resized_height), (0, 0)))
+
+        return tensor_input

samgis_core/prediction_api/sam_onnx_inference.py

@@ -0,0 +1,121 @@
+from numpy import array as np_array, uint8, zeros, ndarray
+
+from samgis_core import app_logger, MODEL_FOLDER
+from samgis_core.prediction_api.sam_onnx2 import SegmentAnythingONNX2
+from samgis_core.utilities.constants import MODEL_ENCODER_NAME, MODEL_DECODER_NAME
+from samgis_core.utilities.type_hints import ListDict, PIL_Image, TupleNdarrayInt, EmbeddingPILDict
+
+
+def get_raster_inference(
+        img: PIL_Image | ndarray, prompt: ListDict, models_instance: SegmentAnythingONNX2, model_name: str
+) -> TupleNdarrayInt:
+    """
+    Get inference output for a given image using a SegmentAnythingONNX model
+
+    Args:
+        img: input PIL Image
+        prompt: list of prompt dict
+        models_instance: SegmentAnythingONNX instance model
+        model_name: model name string
+
+    Returns:
+        raster prediction mask, prediction number
+    """
+    np_img = np_array(img)
+    app_logger.info(f"img type {type(np_img)}, prompt:{prompt}.")
+    app_logger.debug(f"onnxruntime input shape/size (shape if PIL) {np_img.size}.")
+    try:
+        app_logger.debug(f"onnxruntime input shape (NUMPY) {np_img.shape}.")
+    except Exception as e_shape:
+        app_logger.error(f"e_shape:{e_shape}.")
+    app_logger.info(f"instantiated model {model_name}, ENCODER {MODEL_ENCODER_NAME}, "
+                    f"DECODER {MODEL_DECODER_NAME} from {MODEL_FOLDER}: Creating embedding...")
+    embedding = models_instance.encode(np_img)
+    app_logger.debug(f"embedding created, running predict_masks with prompt {prompt}...")
+    return get_raster_inference_using_existing_embedding(embedding, prompt, models_instance)
+
+
+def get_inference_embedding(
+        img: PIL_Image | ndarray, models_instance: SegmentAnythingONNX2, model_name: str, embedding_key: str,
+        embedding_dict: EmbeddingPILDict) -> EmbeddingPILDict:
+    """add an embedding to the embedding dict if needed
+
+    Args:
+        img: input PIL Image
+        models_instance: SegmentAnythingONNX instance model
+        model_name: model name string
+        embedding_key: embedding id
+        embedding_dict: embedding dict object
+
+    Returns:
+        raster dict
+    """
+    if embedding_key in embedding_dict:
+        app_logger.info("found embedding in dict...")
+    if embedding_key not in embedding_dict:
+        np_img = np_array(img)
+        app_logger.info(f"prepare embedding using img type {type(np_img)}.")
+        app_logger.debug(f"onnxruntime input shape/size (shape if PIL) {np_img.size}.")
+        try:
+            app_logger.debug(f"onnxruntime input shape (NUMPY) {np_img.shape}.")
+        except Exception as e_shape:
+            app_logger.error(f"e_shape:{e_shape}.")
+        app_logger.info(f"instantiated model {model_name}, ENCODER {MODEL_ENCODER_NAME}, "
+                        f"DECODER {MODEL_DECODER_NAME} from {MODEL_FOLDER}: Creating embedding...")
+        embedding = models_instance.encode(np_img)
+        embedding_dict[embedding_key] = embedding
+    return embedding_dict
+
+
+def get_raster_inference_using_existing_embedding(
+        embedding: dict, prompt: ListDict, models_instance: SegmentAnythingONNX2) -> TupleNdarrayInt:
+    """
+    Get inference output for a given image using a SegmentAnythingONNX model, using an existing embedding instead of a
+    new ndarray or PIL image
+
+    Args:
+        embedding: dict
+        prompt: list of prompt dict
+        models_instance: SegmentAnythingONNX instance model
+
+    Returns:
+        raster prediction mask, prediction number
+    """
+    app_logger.info(f"using existing embedding of type {type(embedding)}.")
+    inference_out = models_instance.predict_masks(embedding, prompt)
+    len_inference_out = len(inference_out[0, :, :, :])
+    app_logger.info(f"Created {len_inference_out} prediction_masks,"
+                    f"shape:{inference_out.shape}, dtype:{inference_out.dtype}.")
+    mask = zeros((inference_out.shape[2], inference_out.shape[3]), dtype=uint8)
+    for n, m in enumerate(inference_out[0, :, :, :]):
+        app_logger.debug(f"{n}th of prediction_masks shape {inference_out.shape}"
+                         f" => mask shape:{mask.shape}, {mask.dtype}.")
+        mask[m > 0.0] = 255
+    return mask, len_inference_out
+
+
+def get_raster_inference_with_embedding_from_dict(
+        img: PIL_Image | ndarray, prompt: ListDict, models_instance: SegmentAnythingONNX2, model_name: str,
+        embedding_key: str, embedding_dict: dict) -> TupleNdarrayInt:
+    """
+    Get inference output using a SegmentAnythingONNX model, but get the image embedding from the given embedding dict
+     instead of creating a new embedding. This function needs the img argument to update the embedding dict if necessary
+
+    Args:
+        img: input PIL Image
+        prompt: list of prompt dict
+        models_instance: SegmentAnythingONNX instance model
+        model_name: model name string
+        embedding_key: embedding id
+        embedding_dict: embedding images dict
+
+    Returns:
+        raster prediction mask, prediction number
+    """
+    app_logger.info(f"handling embedding using key {embedding_key}.")
+    embedding_dict = get_inference_embedding(img, models_instance, model_name, embedding_key, embedding_dict)
+    app_logger.info(f"getting embedding with key {embedding_key} from dict...")
+    embedding = embedding_dict[embedding_key]
+    n_keys = len(embedding_dict)
+    app_logger.info(f"embedding created ({n_keys} keys in embedding dict), running predict_masks with prompt {prompt}.")
+    return get_raster_inference_using_existing_embedding(embedding, prompt, models_instance)

samgis_core/utilities/constants.py

@@ -1,5 +1,6 @@
 """Project constants"""
-DEFAULT_INPUT_SHAPE = 684, 1024
-MODEL_ENCODER_NAME = "mobile_sam.encoder.onnx"
-MODEL_DECODER_NAME = "sam_vit_h_4b8939.decoder.onnx"
+import os
 
+DEFAULT_INPUT_SHAPE = 684, 1024
+MODEL_ENCODER_NAME = os.getenv("MODEL_ENCODER_NAME", "mobile_sam.encoder.onnx")
+MODEL_DECODER_NAME = os.getenv("MODEL_DECODER_NAME", "mobile_sam.decoder.onnx")

samgis_core/utilities/plot_images.py

@@ -0,0 +1,11 @@
+def helper_imshow_output_expected(img1, img2, title1, title2):
+    from matplotlib import pyplot as plt
+
+    fig, ax = plt.subplot_mosaic([
+        [title1, title2]
+    ], figsize=(15, 10))
+
+    ax[title1].imshow(img1)
+    ax[title2].imshow(img2)
+
+    plt.show()

samgis_core/utilities/type_hints.py

@@ -27,6 +27,9 @@ class ListInt(list[int]): pass
 class TupleInt(tuple[int]): pass
 
 
+TupleInt2 = NewType("TupleInt", tuple[int, int])
+
+
 class TupleNdarrayInt(tuple[ndarray, int]): pass
 
 
@@ -55,4 +58,11 @@ class EmbeddingImage(TypedDict):
     transform_matrix: ndarray
 
 
+class EmbeddingPILImage(TypedDict):
+    image_embedding: ndarray
+    original_size: TupleInt2
+    resized_size: TupleInt2
+
+
 EmbeddingDict = dict[str, EmbeddingImage]
+EmbeddingPILDict = dict[str, EmbeddingPILImage]

samgis_core/utilities/utilities.py

@@ -1,8 +1,11 @@
 """Various utilities (logger, time benchmark, args dump, numerical and stats info)"""
-import numbers
+from copy import deepcopy
+
+from numpy import ndarray, float32
 
 from samgis_core import app_logger
 from samgis_core.utilities.serialize import serialize
+from samgis_core.utilities.type_hints import EmbeddingPILImage, PIL_Image
 
 
 def _prepare_base64_input(sb):
@@ -14,7 +17,7 @@ def _prepare_base64_input(sb):
     raise ValueError("Argument must be string or bytes")
 
 
-def _is_base64(sb: str or bytes):
+def _is_base64(sb: str | bytes):
     import base64
 
     try:
@@ -43,7 +46,7 @@ def base64_decode(s):
     return s
 
 
-def base64_encode(sb: str or bytes) -> bytes:
+def base64_encode(sb: str | bytes) -> bytes:
     """
     Encode input strings or bytes as base64
 
@@ -59,7 +62,7 @@ def base64_encode(sb: str or bytes) -> bytes:
     return base64.b64encode(sb_bytes)
 
 
-def hash_calculate(arr) -> str or bytes:
+def hash_calculate(arr) -> str | bytes:
     """
     Return computed hash from input variable (typically a numpy array).
 
@@ -92,3 +95,26 @@ def hash_calculate(arr) -> str or bytes:
     else:
         raise ValueError(f"variable 'arr':{arr} of type '{type(arr)}' not yet handled.")
     return b64encode(hash_fn.digest())
+
+
+def convert_ndarray_to_pil(pil_image: PIL_Image | ndarray):
+    from PIL import Image
+
+    if isinstance(pil_image, ndarray):
+        pil_image = Image.fromarray(pil_image)
+    return pil_image
+
+
+def apply_coords(coords: ndarray, embedding: EmbeddingPILImage):
+    """
+    Expects a numpy np_array of length 2 in the final dimension. Requires the
+    original image size in (H, W) format.
+    """
+    orig_width, orig_height = embedding["original_size"]
+    resized_width, resized_height = embedding["resized_size"]
+    coords = deepcopy(coords).astype(float)
+
+    coords[..., 0] = coords[..., 0] * (resized_width / orig_width)
+    coords[..., 1] = coords[..., 1] * (resized_height / orig_height)
+
+    return coords.astype(float32)

{samgis_core-1.2.1.dist-info → samgis_core-1.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: samgis_core
-Version: 1.2.1
+Version: 1.2.2
 Summary: SamGIS CORE
 License: MIT
 Author: alessandro trinca tornidor
@@ -14,8 +14,7 @@ Requires-Dist: bson (>=0.5.10,<0.6.0)
 Requires-Dist: loguru (>=0.7.2,<0.8.0)
 Requires-Dist: numpy (==1.25.2) ; python_version >= "3.10" and python_version < "3.11"
 Requires-Dist: numpy (>=1.26,<2.0) ; python_version >= "3.11" and python_version < "3.12"
-Requires-Dist: onnxruntime (==1.16.3)
-Requires-Dist: opencv-python-headless (==4.8.1.78)
+Requires-Dist: onnxruntime (==1.17.3)
 Requires-Dist: pillow (>=10.2.0,<11.0.0)
 Project-URL: Source, https://gitlab.com/aletrn/samgis_core
 Description-Content-Type: text/markdown

samgis_core-1.2.2.dist-info/RECORD

@@ -0,0 +1,15 @@
+samgis_core/__init__.py,sha256=1kFX8G22dxNz23J7uOYl-SMWOe4W1olssc-5zKAVwSc,351
+samgis_core/__version__.py,sha256=x-8LeQkljky-ao1MyT98tSQE1xNvJp3LelqsAdCb9Og,94
+samgis_core/prediction_api/__init__.py,sha256=_jUZhspS26ygiSzBJywfZ4fQf9X7oC8w6oxvhd9S2hQ,57
+samgis_core/prediction_api/sam_onnx2.py,sha256=u0-BgNff0nMzRNk9MOF1-hVMmWJcZr_EFXAp8iF5NBs,7437
+samgis_core/prediction_api/sam_onnx_inference.py,sha256=EQjJy1RLZn0-m66bdY_T1bR-t6PnNu47JJVphEkgPhE,5539
+samgis_core/utilities/__init__.py,sha256=nL9pzdB4SdEF8m5gCbtlVCtdGLg9JjPm-FNxKBsIBZA,32
+samgis_core/utilities/constants.py,sha256=0xBdfGYwCg4O0OXFtTcMVNj-kryjbajcxOZhMVkVP7U,227
+samgis_core/utilities/fastapi_logger.py,sha256=yt8Vj1viyE-Kry1_iy5p4saZvGEEmRfuTEMHYXgnjqk,650
+samgis_core/utilities/plot_images.py,sha256=nY_1KW7x_h218MA9leulEKLkoNhLlyEscLSWs0Az2uE,263
+samgis_core/utilities/serialize.py,sha256=aIjhEoibBpV_gpgOg6LiVxZCWjOkYxlzcboDZLQctJE,2689
+samgis_core/utilities/type_hints.py,sha256=hAMYXpHgMhYguOPbegCQbVCIexWNSn6S-2Q_nilfesQ,1068
+samgis_core/utilities/utilities.py,sha256=tRGp-Iw0PoPf5YHDL6Hx2CEdZhqs03hh3YPe_rmoO-E,3286
+samgis_core-1.2.2.dist-info/METADATA,sha256=ZBYbeGm4TKxV7OVTihqdCXrhRVHSlYkF3zY3KaoFErg,1096
+samgis_core-1.2.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+samgis_core-1.2.2.dist-info/RECORD,,

samgis_core/prediction_api/sam_onnx.py

@@ -1,371 +0,0 @@
-"""
-Define a machine learning model executed by ONNX Runtime (https://ai/)
-for Segment Anything (https://segment-anything.com).
-Modified from https://github.com/vietanhdev/samexporter/
-
-Copyright (c) 2023 Viet Anh Nguyen
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.
-"""
-from copy import deepcopy
-
-from cv2 import INTER_LINEAR, warpAffine
-from numpy import array as np_array, uint8, zeros, ndarray
-from numpy import concatenate, float32, linalg, matmul, ones
-from onnxruntime import get_available_providers, InferenceSession
-
-from samgis_core import app_logger, MODEL_FOLDER
-from samgis_core.utilities.constants import DEFAULT_INPUT_SHAPE, MODEL_ENCODER_NAME, MODEL_DECODER_NAME
-from samgis_core.utilities.type_hints import PIL_Image, ListDict, TupleNdarrayInt, EmbeddingDict, EmbeddingImage
-
-
-class SegmentAnythingONNX:
-    """Segmentation model using SegmentAnything"""
-
-    def __init__(self, encoder_model_path, decoder_model_path) -> None:
-        self.target_size = DEFAULT_INPUT_SHAPE[1]
-        self.input_size = DEFAULT_INPUT_SHAPE
-
-        # Load models
-        providers = get_available_providers()
-
-        # Pop TensorRT Runtime due to crashing issues
-        # TODO: Add back when TensorRT backend is stable
-        providers = [p for p in providers if p != "TensorrtExecutionProvider"]
-
-        if providers:
-            app_logger.info(
-                "Available providers for ONNXRuntime: %s", ", ".join(providers)
-            )
-        else:
-            app_logger.warning("No available providers for ONNXRuntime")
-        self.encoder_session = InferenceSession(
-            encoder_model_path, providers=providers
-        )
-        self.encoder_input_name = self.encoder_session.get_inputs()[0].name
-        self.decoder_session = InferenceSession(
-            decoder_model_path, providers=providers
-        )
-
-    @staticmethod
-    def get_input_points(prompt):
-        """Get input points"""
-        points = []
-        labels = []
-        for mark in prompt:
-            if mark["type"] == "point":
-                points.append(mark["data"])
-                labels.append(mark["label"])
-            elif mark["type"] == "rectangle":
-                points.append([mark["data"][0], mark["data"][1]])  # top left
-                points.append(
-                    [mark["data"][2], mark["data"][3]]
-                )  # bottom right
-                labels.append(2)
-                labels.append(3)
-        points, labels = np_array(points), np_array(labels)
-        return points, labels
-
-    def run_encoder(self, encoder_inputs):
-        """Run encoder"""
-        output = self.encoder_session.run(None, encoder_inputs)
-        image_embedding = output[0]
-        return image_embedding
-
-    @staticmethod
-    def get_preprocess_shape(old_h: int, old_w: int, long_side_length: int):
-        """
-        Compute the output size given input size and target long side length.
-        """
-        scale = long_side_length * 1.0 / max(old_h, old_w)
-        new_h, new_w = old_h * scale, old_w * scale
-        new_w = int(new_w + 0.5)
-        new_h = int(new_h + 0.5)
-        return new_h, new_w
-
-    def apply_coords(self, coords: ndarray, original_size, target_length):
-        """
-        Expects a numpy np_array of length 2 in the final dimension. Requires the
-        original image size in (H, W) format.
-        """
-        old_h, old_w = original_size
-        new_h, new_w = self.get_preprocess_shape(
-            original_size[0], original_size[1], target_length
-        )
-        coords = deepcopy(coords).astype(float)
-        coords[..., 0] = coords[..., 0] * (new_w / old_w)
-        coords[..., 1] = coords[..., 1] * (new_h / old_h)
-        return coords
-
-    def run_decoder(
-            self, image_embedding, original_size, transform_matrix, prompt
-    ):
-        """Run decoder"""
-        input_points, input_labels = self.get_input_points(prompt)
-
-        # Add a batch index, concatenate a padding point, and transform.
-        onnx_coord = concatenate(
-            [input_points, np_array([[0.0, 0.0]])], axis=0
-        )[None, :, :]
-        onnx_label = concatenate([input_labels, np_array([-1])], axis=0)[
-                     None, :
-                     ].astype(float32)
-        onnx_coord = self.apply_coords(
-            onnx_coord, self.input_size, self.target_size
-        ).astype(float32)
-
-        # Apply the transformation matrix to the coordinates.
-        onnx_coord = concatenate(
-            [
-                onnx_coord,
-                ones((1, onnx_coord.shape[1], 1), dtype=float32),
-            ],
-            axis=2,
-        )
-        onnx_coord = matmul(onnx_coord, transform_matrix.T)
-        onnx_coord = onnx_coord[:, :, :2].astype(float32)
-
-        # Create an empty mask input and an indicator for no mask.
-        onnx_mask_input = zeros((1, 1, 256, 256), dtype=float32)
-        onnx_has_mask_input = zeros(1, dtype=float32)
-
-        decoder_inputs = {
-            "image_embeddings": image_embedding,
-            "point_coords": onnx_coord,
-            "point_labels": onnx_label,
-            "mask_input": onnx_mask_input,
-            "has_mask_input": onnx_has_mask_input,
-            "orig_im_size": np_array(self.input_size, dtype=float32),
-        }
-        masks, _, _ = self.decoder_session.run(None, decoder_inputs)
-
-        # Transform the masks back to the original image size.
-        inv_transform_matrix = linalg.inv(transform_matrix)
-        transformed_masks = self.transform_masks(
-            masks, original_size, inv_transform_matrix
-        )
-
-        return transformed_masks
-
-    @staticmethod
-    def transform_masks(masks, original_size, transform_matrix):
-        """Transform masks
-        Transform the masks back to the original image size.
-        """
-        output_masks = []
-        for batch in range(masks.shape[0]):
-            batch_masks = []
-            for mask_id in range(masks.shape[1]):
-                mask = masks[batch, mask_id]
-                try:
-                    try:
-                        app_logger.debug(f"mask_shape transform_masks:{mask.shape}, dtype:{mask.dtype}.")
-                    except Exception as e_mask_shape_transform_masks:
-                        app_logger.error(f"e_mask_shape_transform_masks:{e_mask_shape_transform_masks}.")
-                    mask = warpAffine(
-                        mask,
-                        transform_matrix[:2],
-                        (original_size[1], original_size[0]),
-                        flags=INTER_LINEAR,
-                    )
-                except Exception as e_warp_affine1:
-                    app_logger.error(f"e_warp_affine1 mask shape:{mask.shape}, dtype:{mask.dtype}.")
-                    app_logger.error(
-                        f"e_warp_affine1 transform_matrix:{transform_matrix}, [:2] {transform_matrix[:2]}.")
-                    app_logger.error(f"e_warp_affine1 original_size:{original_size}.")
-                    raise e_warp_affine1
-                batch_masks.append(mask)
-            output_masks.append(batch_masks)
-        return np_array(output_masks)
-
-    def encode(self, cv_image: ndarray) -> EmbeddingImage:
-        """
-        Calculate embedding and metadata for a single image.
-
-        Args:
-            cv_image: input image to embed
-
-        Returns:
-            embedding image dict useful to store and cache image embeddings
-        """
-        original_size = cv_image.shape[:2]
-
-        # Calculate a transformation matrix to convert to self.input_size
-        scale_x = self.input_size[1] / cv_image.shape[1]
-        scale_y = self.input_size[0] / cv_image.shape[0]
-        scale = min(scale_x, scale_y)
-        transform_matrix = np_array(
-            [
-                [scale, 0, 0],
-                [0, scale, 0],
-                [0, 0, 1],
-            ]
-        )
-        try:
-            cv_image = warpAffine(
-                cv_image,
-                transform_matrix[:2],
-                (self.input_size[1], self.input_size[0]),
-                flags=INTER_LINEAR,
-            )
-        except Exception as e_warp_affine2:
-            app_logger.error(f"e_warp_affine2:{e_warp_affine2}.")
-            np_cv_image = np_array(cv_image)
-            app_logger.error(f"e_warp_affine2 cv_image shape:{np_cv_image.shape}, dtype:{np_cv_image.dtype}.")
-            app_logger.error(f"e_warp_affine2 transform_matrix:{transform_matrix}, [:2] {transform_matrix[:2]}")
-            app_logger.error(f"e_warp_affine2 self.input_size:{self.input_size}.")
-            raise e_warp_affine2
-
-        encoder_inputs = {
-            self.encoder_input_name: cv_image.astype(float32),
-        }
-        image_embedding = self.run_encoder(encoder_inputs)
-        return {
-            "image_embedding": image_embedding,
-            "original_size": original_size,
-            "transform_matrix": transform_matrix,
-        }
-
-    def predict_masks(self, embedding, prompt):
-        """
-        Predict masks for a single image.
-        """
-        masks = self.run_decoder(
-            embedding["image_embedding"],
-            embedding["original_size"],
-            embedding["transform_matrix"],
-            prompt,
-        )
-
-        return masks
-
-
-def get_raster_inference(
-        img: PIL_Image or ndarray, prompt: ListDict, models_instance: SegmentAnythingONNX, model_name: str
-) -> TupleNdarrayInt:
-    """
-    Get inference output for a given image using a SegmentAnythingONNX model
-
-    Args:
-        img: input PIL Image
-        prompt: list of prompt dict
-        models_instance: SegmentAnythingONNX instance model
-        model_name: model name string
-
-    Returns:
-        raster prediction mask, prediction number
-    """
-    np_img = np_array(img)
-    app_logger.info(f"img type {type(np_img)}, prompt:{prompt}.")
-    app_logger.debug(f"onnxruntime input shape/size (shape if PIL) {np_img.size}.")
-    try:
-        app_logger.debug(f"onnxruntime input shape (NUMPY) {np_img.shape}.")
-    except Exception as e_shape:
-        app_logger.error(f"e_shape:{e_shape}.")
-    app_logger.info(f"instantiated model {model_name}, ENCODER {MODEL_ENCODER_NAME}, "
-                    f"DECODER {MODEL_DECODER_NAME} from {MODEL_FOLDER}: Creating embedding...")
-    embedding = models_instance.encode(np_img)
-    app_logger.debug(f"embedding created, running predict_masks with prompt {prompt}...")
-    return get_raster_inference_using_existing_embedding(embedding, prompt, models_instance)
-
-
-def get_inference_embedding(
-        img: PIL_Image or ndarray, models_instance: SegmentAnythingONNX, model_name: str, embedding_key: str,
-        embedding_dict: EmbeddingDict) -> EmbeddingDict:
-    """add an embedding to the embedding dict if needed
-
-    Args:
-        img: input PIL Image
-        models_instance: SegmentAnythingONNX instance model
-        model_name: model name string
-        embedding_key: embedding id
-        embedding_dict: embedding dict object
-
-    Returns:
-        raster dict
-    """
-    if embedding_key in embedding_dict:
-        app_logger.info("found embedding in dict...")
-    if embedding_key not in embedding_dict:
-        np_img = np_array(img)
-        app_logger.info(f"prepare embedding using img type {type(np_img)}.")
-        app_logger.debug(f"onnxruntime input shape/size (shape if PIL) {np_img.size}.")
-        try:
-            app_logger.debug(f"onnxruntime input shape (NUMPY) {np_img.shape}.")
-        except Exception as e_shape:
-            app_logger.error(f"e_shape:{e_shape}.")
-        app_logger.info(f"instantiated model {model_name}, ENCODER {MODEL_ENCODER_NAME}, "
-                        f"DECODER {MODEL_DECODER_NAME} from {MODEL_FOLDER}: Creating embedding...")
-        embedding = models_instance.encode(np_img)
-        embedding_dict[embedding_key] = embedding
-    return embedding_dict
-
-
-def get_raster_inference_using_existing_embedding(
-        embedding: dict, prompt: ListDict, models_instance: SegmentAnythingONNX) -> TupleNdarrayInt:
-    """
-    Get inference output for a given image using a SegmentAnythingONNX model, using an existing embedding instead of a
-    new ndarray or PIL image
-
-    Args:
-        embedding: dict
-        prompt: list of prompt dict
-        models_instance: SegmentAnythingONNX instance model
-
-    Returns:
-        raster prediction mask, prediction number
-    """
-    app_logger.info(f"using existing embedding of type {type(embedding)}.")
-    inference_out = models_instance.predict_masks(embedding, prompt)
-    len_inference_out = len(inference_out[0, :, :, :])
-    app_logger.info(f"Created {len_inference_out} prediction_masks,"
-                    f"shape:{inference_out.shape}, dtype:{inference_out.dtype}.")
-    mask = zeros((inference_out.shape[2], inference_out.shape[3]), dtype=uint8)
-    for n, m in enumerate(inference_out[0, :, :, :]):
-        app_logger.debug(f"{n}th of prediction_masks shape {inference_out.shape}"
-                         f" => mask shape:{mask.shape}, {mask.dtype}.")
-        mask[m > 0.0] = 255
-    return mask, len_inference_out
-
-
-def get_raster_inference_with_embedding_from_dict(
-        img: PIL_Image or ndarray, prompt: ListDict, models_instance: SegmentAnythingONNX, model_name: str,
-        embedding_key: str, embedding_dict: dict) -> TupleNdarrayInt:
-    """
-    Get inference output using a SegmentAnythingONNX model, but get the image embedding from the given embedding dict
-     instead of creating a new embedding. This function needs the img argument to update the embedding dict if necessary
-
-    Args:
-        img: input PIL Image
-        prompt: list of prompt dict
-        models_instance: SegmentAnythingONNX instance model
-        model_name: model name string
-        embedding_key: embedding id
-        embedding_dict: embedding images dict
-
-    Returns:
-        raster prediction mask, prediction number
-    """
-    app_logger.info(f"handling embedding using key {embedding_key}.")
-    embedding_dict = get_inference_embedding(img, models_instance, model_name, embedding_key, embedding_dict)
-    app_logger.info(f"getting embedding with key {embedding_key} from dict...")
-    embedding = embedding_dict[embedding_key]
-    n_keys = len(embedding_dict)
-    app_logger.info(f"embedding created ({n_keys} keys in embedding dict), running predict_masks with prompt {prompt}.")
-    return get_raster_inference_using_existing_embedding(embedding, prompt, models_instance)

samgis_core-1.2.1.dist-info/RECORD

@@ -1,13 +0,0 @@
-samgis_core/__init__.py,sha256=1kFX8G22dxNz23J7uOYl-SMWOe4W1olssc-5zKAVwSc,351
-samgis_core/__version__.py,sha256=x-8LeQkljky-ao1MyT98tSQE1xNvJp3LelqsAdCb9Og,94
-samgis_core/prediction_api/__init__.py,sha256=_jUZhspS26ygiSzBJywfZ4fQf9X7oC8w6oxvhd9S2hQ,57
-samgis_core/prediction_api/sam_onnx.py,sha256=YEX6UOFKcqPiuF_SCPAcK_gCmxoj4lvTLNNDul2htZ8,15320
-samgis_core/utilities/__init__.py,sha256=nL9pzdB4SdEF8m5gCbtlVCtdGLg9JjPm-FNxKBsIBZA,32
-samgis_core/utilities/constants.py,sha256=645W57jrUNbnujgUwCietfr-rECENDXLGmHeD2YoSwg,157
-samgis_core/utilities/fastapi_logger.py,sha256=yt8Vj1viyE-Kry1_iy5p4saZvGEEmRfuTEMHYXgnjqk,650
-samgis_core/utilities/serialize.py,sha256=aIjhEoibBpV_gpgOg6LiVxZCWjOkYxlzcboDZLQctJE,2689
-samgis_core/utilities/type_hints.py,sha256=iDkWiVB2B0W_GaRhdpypEra2FvQtwC5lNLVZr7a4yhU,845
-samgis_core/utilities/utilities.py,sha256=fP3cnxIYULYoCBft54EAivUR_fMcCDz2Z9AMpAO0zZ0,2434
-samgis_core-1.2.1.dist-info/METADATA,sha256=L0FMNdjcgVtJnGr3j78buKJlry9o340vjIWL0b_TUrQ,1147
-samgis_core-1.2.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-samgis_core-1.2.1.dist-info/RECORD,,