[feat] take advantage of re-usable image embeddings in SAM model

pyproject.toml

@@ -11,15 +11,15 @@ bson = "^0.5.10"
 contextily = "^1.5.2"
 geopandas = "^0.14.3"
 loguru = "^0.7.2"
-numpy = "^1.26.4"
+numpy = "~1.25.2"
 onnxruntime = "1.16.3"
 opencv-python-headless = "^4.8.1.78"
 pillow = "^10.2.0"
-python = "^3.11"
+python = "~3.10"
 python-dotenv = "^1.0.1"
 rasterio = "^1.3.9"
 requests = "^2.31.0"
-samgis-core = "^1.0.3"
+samgis-core = "^1.1.1"
 
 [tool.poetry.group.aws_lambda]
 optional = true

samgis/io/wrappers_helpers.py

@@ -200,3 +200,25 @@ def get_url_tile(source_type: str):
 
 def check_source_type_is_terrain(source: str | TileProvider):
     return isinstance(source, TileProvider) and source.name in list(XYZTerrainProvidersNames)
+
+
+def get_source_name(source: str | TileProvider) -> str | bool:
+    try:
+        match source.lower():
+            case XYZDefaultProvidersNames.DEFAULT_TILES_NAME_SHORT:
+                source_output = providers.query_name(XYZDefaultProvidersNames.DEFAULT_TILES_NAME)
+            case _:
+                source_output = providers.query_name(source)
+        if isinstance(source_output, str):
+            return source_output
+        try:
+            source_dict = dict(source_output)
+            app_logger.info(f"source_dict:{type(source_dict)}, {'name' in source_dict}, source_dict:{source_dict}.")
+            return source_dict["name"]
+        except KeyError as ke:
+            app_logger.error(f"ke:{ke}.")
+    except ValueError as ve:
+        app_logger.info(f"source name::{source}, ve:{ve}.")
+    app_logger.info(f"source name::{source}.")
+
+    return False

samgis/prediction_api/predictors.py

@@ -6,12 +6,13 @@ from samgis.io.tms2geotiff import download_extent
 from samgis.io.wrappers_helpers import check_source_type_is_terrain
 from samgis.utilities.constants import DEFAULT_URL_TILES, SLOPE_CELLSIZE
 from samgis_core.prediction_api.sam_onnx import SegmentAnythingONNX
-from samgis_core.prediction_api.sam_onnx import get_raster_inference
+from samgis_core.prediction_api.sam_onnx import get_raster_inference, get_raster_inference_with_embedding_from_dict
 from samgis_core.utilities.constants import MODEL_ENCODER_NAME, MODEL_DECODER_NAME, DEFAULT_INPUT_SHAPE
 from samgis_core.utilities.type_hints import llist_float, dict_str_int, list_dict
 
 
 models_dict = {"fastsam": {"instance": None}}
+embedding_dict = {}
 
 
 def samexporter_predict(
@@ -19,7 +20,8 @@ def samexporter_predict(
         prompt: list_dict,
         zoom: float,
         model_name: str = "fastsam",
-        source: str = DEFAULT_URL_TILES
+        source: str = DEFAULT_URL_TILES,
+        source_name: str = None
 ) -> dict_str_int:
     """
     Return predictions as a geojson from a geo-referenced image using the given input prompt.
@@ -34,7 +36,8 @@ def samexporter_predict(
         prompt: machine learning input prompt
         zoom: Level of detail
         model_name: machine learning model name
-        source: xyz
+        source: xyz tile provider object
+        source_name: name of tile provider
 
     Returns:
         Affine transform
@@ -62,9 +65,12 @@ def samexporter_predict(
 
     app_logger.info(
         f"img type {type(img)} with shape/size:{img.size}, transform type: {type(transform)}, transform:{transform}.")
-
-    mask, n_predictions = get_raster_inference(img, prompt, models_instance, model_name)
-    app_logger.info(f"created {n_predictions} masks, preparing conversion to geojson...")
+    app_logger.info(f"source_name:{source_name}, source_name type:{type(source_name)}.")
+    embedding_key = f"{source_name}_z{zoom}_w{pt1[1]},s{pt1[0]},e{pt0[1]},n{pt0[0]}"
+    mask, n_predictions = get_raster_inference_with_embedding_from_dict(
+        img, prompt, models_instance, model_name, embedding_key, embedding_dict)
+    app_logger.info(f"created {n_predictions} masks, type {type(mask)}, size {mask.size}: preparing geojson conversion")
+    app_logger.info(f"mask shape:{mask.shape}.")
     return {
         "n_predictions": n_predictions,
         **get_vectorized_raster_as_geojson(mask, transform)

wrappers/fastapi_wrapper.py

@@ -8,13 +8,14 @@ from fastapi.staticfiles import StaticFiles
 from pydantic import ValidationError
 
 from samgis import PROJECT_ROOT_FOLDER
-from samgis.io.wrappers_helpers import get_parsed_bbox_points
+from samgis.io.wrappers_helpers import get_parsed_bbox_points, get_source_name
 from samgis.utilities.type_hints import ApiRequestBody
 from samgis_core.utilities.fastapi_logger import setup_logging
 from samgis.prediction_api.predictors import samexporter_predict
 
 
 app_logger = setup_logging(debug=True)
+app_logger.info(f"PROJECT_ROOT_FOLDER:{PROJECT_ROOT_FOLDER}.")
 app = FastAPI()
 
 
@@ -68,9 +69,11 @@ def infer_samgis(request_input: ApiRequestBody) -> JSONResponse:
         body_request = get_parsed_bbox_points(request_input)
         app_logger.info(f"body_request:{body_request}.")
         try:
+            source_name = get_source_name(request_input.source_type)
+            app_logger.info(f"source_name = {source_name}.")
             output = samexporter_predict(
                 bbox=body_request["bbox"], prompt=body_request["prompt"], zoom=body_request["zoom"],
-                source=body_request["source"]
+                source=body_request["source"], source_name=source_name
             )
             duration_run = time.time() - time_start_run
             app_logger.info(f"duration_run:{duration_run}.")