samgis_core 1.0.8__py3-none-any.whl → 1.1.0__py3-none-any.whl

samgis_core/prediction_api/sam_onnx.py

@@ -32,7 +32,7 @@ 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, list_dict, tuple_ndarr_int
+from samgis_core.utilities.type_hints import PIL_Image, list_dict, tuple_ndarr_int, EmbeddingDict, EmbeddingImage
 
 
 class SegmentAnythingONNX:
@@ -114,7 +114,7 @@ class SegmentAnythingONNX:
         return coords
 
     def run_decoder(
-        self, image_embedding, original_size, transform_matrix, prompt
+            self, image_embedding, original_size, transform_matrix, prompt
     ):
         """Run decoder"""
         input_points, input_labels = self.get_input_points(prompt)
@@ -124,8 +124,8 @@ class SegmentAnythingONNX:
             [input_points, np_array([[0.0, 0.0]])], axis=0
         )[None, :, :]
         onnx_label = concatenate([input_labels, np_array([-1])], axis=0)[
-            None, :
-        ].astype(float32)
+                     None, :
+                     ].astype(float32)
         onnx_coord = self.apply_coords(
             onnx_coord, self.input_size, self.target_size
         ).astype(float32)
@@ -186,14 +186,15 @@ class SegmentAnythingONNX:
                     )
                 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 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):
+    def encode(self, cv_image: ndarray) -> EmbeddingImage:
         """
         Calculate embedding and metadata for a single image.
         """
@@ -251,7 +252,7 @@ class SegmentAnythingONNX:
 
 def get_raster_inference(
         img: PIL_Image or ndarray, prompt: list_dict, models_instance: SegmentAnythingONNX, model_name: str
-     ) -> tuple_ndarr_int:
+) -> tuple_ndarr_int:
     """
     Get inference output for a given image using a SegmentAnythingONNX model
 
@@ -275,6 +276,56 @@ def get_raster_inference(
                     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: list_dict, models_instance: SegmentAnythingONNX) -> tuple_ndarr_int:
+    """
+    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,"
@@ -285,3 +336,30 @@ def get_raster_inference(
                          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: list_dict, models_instance: SegmentAnythingONNX, model_name: str,
+        embedding_key: str, embedding_dict: dict) -> tuple_ndarr_int:
+    """
+    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 dict object
+
+    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/fastapi_logger.py

@@ -1,8 +1,10 @@
 import loguru
 
 
-def setup_logging(debug: bool = False, formatter: str = "{time} - {level} - ({extra[request_id]}) {message} "
-                  ) -> loguru.logger:
+format_string = "{time} - {level} - {file} - {function} - ({extra[request_id]}) {message} "
+
+
+def setup_logging(debug: bool = False, formatter: str = format_string) -> loguru.logger:
     """
     Create a logging instance with log string formatter.
 

samgis_core/utilities/type_hints.py

@@ -1,7 +1,11 @@
 """custom type hints"""
+from enum import Enum
+from typing import TypedDict
+
 from PIL.Image import Image
 from numpy import ndarray
 
+
 dict_str_int = dict[str, int]
 dict_str = dict[str]
 dict_str_any = dict[str, any]
@@ -14,3 +18,16 @@ llist_float = list[list_float]
 tuple_float = tuple[float]
 tuple_float_any = tuple[float, any]
 PIL_Image = Image
+
+
+class StrEnum(str, Enum):
+    pass
+
+
+class EmbeddingImage(TypedDict):
+    image_embedding: ndarray
+    original_size: tuple_int
+    transform_matrix: ndarray
+
+
+EmbeddingDict = dict[str, EmbeddingImage]

{samgis_core-1.0.8.dist-info → samgis_core-1.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: samgis_core
-Version: 1.0.8
+Version: 1.1.0
 Summary: SamGIS CORE
 License: MIT
 Author: alessandro trinca tornidor

{samgis_core-1.0.8.dist-info → samgis_core-1.1.0.dist-info}/RECORD

@@ -1,13 +1,13 @@
 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=_sCnu2ScJzoi7_XrocbdNCZ5EjetkgBTrGannhzqFfU,11674
+samgis_core/prediction_api/sam_onnx.py,sha256=G5jJUyidLC_-9VrVmQRUHQpKQHW-vnv_EC7ppNap9VI,15172
 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=c-UTWPbshcbma7EnZd98mN8bxz3GVdBsEe3HKm7C5cQ,615
+samgis_core/utilities/fastapi_logger.py,sha256=yt8Vj1viyE-Kry1_iy5p4saZvGEEmRfuTEMHYXgnjqk,650
 samgis_core/utilities/serialize.py,sha256=iWi_m_7vL7Hmr3Lb5yASXU5DtoUU18M_JdwiRSOs9sI,2749
-samgis_core/utilities/type_hints.py,sha256=YoP7T9ftYKEq0c8xlsahtYewmu_jjUY2-EFBszauBlI,402
+samgis_core/utilities/type_hints.py,sha256=R-xfFbiGz12tuHqChlVI9GguR9Q7pxceFkCi_OmscM4,659
 samgis_core/utilities/utilities.py,sha256=fP3cnxIYULYoCBft54EAivUR_fMcCDz2Z9AMpAO0zZ0,2434
-samgis_core-1.0.8.dist-info/METADATA,sha256=Mq6aSqoUg_PcoI-U4_9LcCog3Qws72jgm6GQ4jzjpBI,892
-samgis_core-1.0.8.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-samgis_core-1.0.8.dist-info/RECORD,,
+samgis_core-1.1.0.dist-info/METADATA,sha256=Zj7E1tpVyMtVbYH4uKYCTNqdYPrjuzgabs878EkNONw,892
+samgis_core-1.1.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+samgis_core-1.1.0.dist-info/RECORD,,