[debug] first test on lambda for fastsam prediction

README.md

@@ -1,5 +1,16 @@
 # Segment Geospatial
 
+## todo
+
+1. export output to mask: OK local, OK aws lambda
+2. resolve model paths: OK local
+3. inference: 
+4. from mask to json (rasterio + geopandas, check for re-projection to EPSG_4326)
+5. check mandatory dependencies
+6. check for alternative python interpreters
+
+## Build instructions
+
 Build the docker image:
 
 ```bash

dockerfiles/dockerfile-lambda-gdal-runner

@@ -9,12 +9,15 @@ RUN echo "ENV RIE: $RIE ..."
 
 # Set working directory to function root directory
 WORKDIR ${LAMBDA_TASK_ROOT}
-COPY requirements.txt ${LAMBDA_TASK_ROOT}/requirements.txt
+COPY requirements_dev.txt ${LAMBDA_TASK_ROOT}/requirements_dev.txt
 
-RUN apt update && apt install -y curl python3-pip
+# avoid segment-geospatial exception caused by missing libGL.so.1 library
+RUN apt update && apt install -y libgl1 curl python3-pip
+RUN ls -ld /usr/lib/*linux-gnu/libGL.so* || echo "libGL.so* not found..."
 RUN which python
 RUN python --version
-RUN python -m pip install pillow awslambdaric aws-lambda-powertools httpx jmespath
+RUN python -m pip install -r ${LAMBDA_TASK_ROOT}/requirements_dev.txt --target ${LAMBDA_TASK_ROOT}
+# RUN python -m pip install pillow awslambdaric aws-lambda-powertools httpx jmespath --target ${LAMBDA_TASK_ROOT}
 
 RUN curl -Lo /usr/local/bin/aws-lambda-rie ${RIE}
 RUN chmod +x /usr/local/bin/aws-lambda-rie

dockerfiles/dockerfile-lambda-samgeo-api

@@ -7,6 +7,7 @@ ARG PYTHONPATH="${LAMBDA_TASK_ROOT}:${PYTHONPATH}:/usr/local/lib/python3/dist-pa
 # Set working directory to function root directory
 WORKDIR ${LAMBDA_TASK_ROOT}
 COPY ./src ${LAMBDA_TASK_ROOT}/src
+COPY ./models ${LAMBDA_TASK_ROOT}/models
 
 RUN ls -l /usr/bin/which
 RUN /usr/bin/which python
@@ -16,8 +17,11 @@ RUN echo "PATH: ${PATH}."
 RUN echo "LAMBDA_TASK_ROOT: ${LAMBDA_TASK_ROOT}."
 RUN ls -l ${LAMBDA_TASK_ROOT}
 RUN ls -ld ${LAMBDA_TASK_ROOT}
+RUN ls -l ${LAMBDA_TASK_ROOT}/models
 RUN python -c "import sys; print(sys.path)"
 RUN python -c "import osgeo"
+RUN python -c "import cv2"
+RUN python -c "import onnxruntime"
 # RUN python -c "import rasterio"
 RUN python -c "import awslambdaric"
 RUN python -m pip list

requirements.txt

@@ -1,8 +1,9 @@
 aws-lambda-powertools
 awslambdaric
 bson
-geojson-pydantic
+httpx
 jmespath
-python-dotenv
-segment-anything-fast
-segment-geospatial
+numpy
+onnxruntime
+opencv-python
+pillow

requirements_dev.txt

@@ -1,7 +1,9 @@
+aws-lambda-powertools
 awslambdaric
-aws_lambda_powertools
-fastjsonschema
-geojson-pydantic
+bson
+httpx
 jmespath
-pydantic
-requests
+numpy
+onnxruntime
+opencv-python
+pillow

src/__init__.py

@@ -1,4 +1,8 @@
 from aws_lambda_powertools import Logger
+import os
+from pathlib import Path
 
 
+PROJECT_ROOT_FOLDER = Path(globals().get("__file__", "./_")).absolute().parent.parent
+MODEL_FOLDER = Path(os.path.join(PROJECT_ROOT_FOLDER, "models"))
 app_logger = Logger()

src/app.py

@@ -1,16 +1,16 @@
 import json
 import time
 from http import HTTPStatus
-
+from typing import Dict
 from aws_lambda_powertools.event_handler import content_types
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 from src import app_logger
-from src.io.tms2geotiff import download_extent
-from src.utilities.constants import CUSTOM_RESPONSE_MESSAGES, DEFAULT_TMS
+from src.prediction_api.predictors import samexporter_predict
+from src.utilities.constants import CUSTOM_RESPONSE_MESSAGES
 
 
-def get_response(status: int, start_time: float, request_id: str, response_body = None) -> str:
+def get_response(status: int, start_time: float, request_id: str, response_body: Dict = None) -> str:
     """
     Return a response for frontend clients.
 
@@ -51,10 +51,20 @@ def lambda_handler(event: dict, context: LambdaContext):
         app_logger.info(f"context:{context}...")
 
         try:
+            """
+            img, matrix = download_extent(DEFAULT_TMS, pt0[0], pt0[1], pt1[0], pt1[1], 6)
+            model_path = Path(MODEL_FOLDER) / "mobile_sam.encoder.onnx"
+            model_path_isfile = model_path.is_file()
+            model_path_stats = model_path.stat()
+            app_logger.info(f"model_path:{model_path_isfile}, {model_path_stats}.")
+            """
             pt0 = 45.699, 127.1
             pt1 = 30.1, 148.492
-            img, matrix = download_extent(DEFAULT_TMS, pt0[0], pt0[1], pt1[0], pt1[1], 6)
-            body_response = {"geojson": {"img_size": img.size}}
+            bbox = [pt0, pt1]
+            zoom = 6
+            prompt = [{"type": "rectangle", "data": [400, 460, 524, 628]}]
+            body_response = {"geojson": samexporter_predict(bbox, prompt, zoom)}
+            app_logger.info(f"body_response::output:{body_response}.")
             response = get_response(HTTPStatus.OK.value, start_time, context.aws_request_id, body_response)
         except Exception as ve:
             app_logger.error(f"validation error:{ve}.")

src/prediction_api/predictors.py

@@ -0,0 +1,82 @@
+# Press the green button in the gutter to run the script.
+import os
+import numpy as np
+
+from src import app_logger, MODEL_FOLDER
+from src.io.tms2geotiff import download_extent
+from src.prediction_api.sam_onnx import SegmentAnythingONNX
+from src.utilities.constants import ROOT, MODEL_ENCODER_NAME, ZOOM, SOURCE_TYPE, DEFAULT_TMS, MODEL_DECODER_NAME
+from src.utilities.serialize import serialize
+from src.utilities.type_hints import input_float_tuples
+from src.utilities.utilities import get_system_info
+
+
+def zip_arrays(arr1, arr2):
+    arr1_list = arr1.tolist()
+    arr2_list = arr2.tolist()
+    # return {serialize(k): serialize(v) for k, v in zip(arr1_list, arr2_list)}
+    d = {}
+    for n1, n2 in enumerate(zip(arr1_list, arr2_list)):
+        app_logger.info(f"n1:{n1}, type {type(n1)}, n2:{n2}, type {type(n2)}.")
+        n1f = str(n1)
+        n2f = str(n2)
+        app_logger.info(f"n1:{n1}=>{n1f}, n2:{n2}=>{n2f}.")
+        d[n1f] = n2f
+    app_logger.info(f"zipped dict:{d}.")
+    return d
+
+
+def samexporter_predict(bbox: input_float_tuples, prompt: list[dict], zoom: float = ZOOM) -> dict:
+    import tempfile
+
+    try:
+        os.environ['MPLCONFIGDIR'] = ROOT
+        get_system_info()
+    except Exception as e:
+        app_logger.error(f"Error while setting 'MPLCONFIGDIR':{e}.")
+
+    with tempfile.NamedTemporaryFile(prefix=f"{SOURCE_TYPE}_", suffix=".tif", dir=ROOT) as image_input_tmp:
+        for coord in bbox:
+            app_logger.info(f"bbox coord:{coord}, type:{type(coord)}.")
+        app_logger.info(f"start download_extent using bbox:{bbox}, type:{type(bbox)}, download image...")
+
+        pt0 = bbox[0]
+        pt1 = bbox[1]
+        img, matrix = download_extent(DEFAULT_TMS, pt0[0], pt0[1], pt1[0], pt1[1], zoom)
+
+        app_logger.info(f"img type {type(img)}, matrix type {type(matrix)}.")
+        np_img = np.array(img)
+        app_logger.info(f"np_img type {type(np_img)}.")
+        app_logger.info(f"np_img dtype {np_img.dtype}, shape {np_img.shape}.")
+        app_logger.info(f"geotiff created with size/shape {img.size} and transform matrix {str(matrix)}, start to initialize SamGeo instance:")
+        app_logger.info(f"use ENCODER model {MODEL_ENCODER_NAME} from {MODEL_FOLDER})...")
+        app_logger.info(f"use DECODER model {MODEL_DECODER_NAME} from {MODEL_FOLDER})...")
+
+        model = SegmentAnythingONNX(
+            encoder_model_path=MODEL_FOLDER / MODEL_ENCODER_NAME,
+            decoder_model_path=MODEL_FOLDER / MODEL_DECODER_NAME
+        )
+        app_logger.info(f"model instantiated, creating embedding...")
+        embedding = model.encode(np_img)
+        app_logger.info(f"embedding created, running predict_masks...")
+        prediction_masks = model.predict_masks(embedding, prompt)
+        app_logger.info(f"predict_masks terminated")
+        app_logger.info(f"prediction masks shape:{prediction_masks.shape}, {prediction_masks.dtype}.")
+
+        mask = np.zeros((prediction_masks.shape[2], prediction_masks.shape[3]), dtype=np.uint8)
+        for m in prediction_masks[0, :, :, :]:
+            mask[m > 0.0] = 255
+
+        mask_unique_values, mask_unique_values_count = serialize(np.unique(mask, return_counts=True))
+        app_logger.info(f"mask_unique_values:{mask_unique_values}.")
+        app_logger.info(f"mask_unique_values_count:{mask_unique_values_count}.")
+
+        output = {
+            "img_size": serialize(img.size),
+            "mask_unique_values_count": zip_arrays(mask_unique_values, mask_unique_values_count),
+            "masks_dtype": serialize(prediction_masks.dtype),
+            "masks_shape": serialize(prediction_masks.shape),
+            "matrix": serialize(matrix)
+        }
+        app_logger.info(f"output:{output}.")
+        return output

src/prediction_api/sam_onnx.py

@@ -0,0 +1,203 @@
+from copy import deepcopy
+
+import cv2
+import numpy as np
+import onnxruntime
+
+from src import app_logger
+
+
+class SegmentAnythingONNX:
+    """Segmentation model using SegmentAnything"""
+
+    def __init__(self, encoder_model_path, decoder_model_path) -> None:
+        self.target_size = 1024
+        self.input_size = (684, 1024)
+
+        # Load models
+        providers = onnxruntime.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 = onnxruntime.InferenceSession(
+            encoder_model_path, providers=providers
+        )
+        self.encoder_input_name = self.encoder_session.get_inputs()[0].name
+        self.decoder_session = onnxruntime.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: np.ndarray, original_size, target_length):
+        """
+        Expects a numpy 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 = np.concatenate(
+            [input_points, np.array([[0.0, 0.0]])], axis=0
+        )[None, :, :]
+        onnx_label = np.concatenate([input_labels, np.array([-1])], axis=0)[
+            None, :
+        ].astype(np.float32)
+        onnx_coord = self.apply_coords(
+            onnx_coord, self.input_size, self.target_size
+        ).astype(np.float32)
+
+        # Apply the transformation matrix to the coordinates.
+        onnx_coord = np.concatenate(
+            [
+                onnx_coord,
+                np.ones((1, onnx_coord.shape[1], 1), dtype=np.float32),
+            ],
+            axis=2,
+        )
+        onnx_coord = np.matmul(onnx_coord, transform_matrix.T)
+        onnx_coord = onnx_coord[:, :, :2].astype(np.float32)
+
+        # Create an empty mask input and an indicator for no mask.
+        onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)
+        onnx_has_mask_input = np.zeros(1, dtype=np.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=np.float32),
+        }
+        masks, _, _ = self.decoder_session.run(None, decoder_inputs)
+
+        # Transform the masks back to the original image size.
+        inv_transform_matrix = np.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]
+                mask = cv2.warpAffine(
+                    mask,
+                    transform_matrix[:2],
+                    (original_size[1], original_size[0]),
+                    flags=cv2.INTER_LINEAR,
+                )
+                batch_masks.append(mask)
+            output_masks.append(batch_masks)
+        return np.array(output_masks)
+
+    def encode(self, cv_image):
+        """
+        Calculate embedding and metadata for a single image.
+        """
+        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],
+            ]
+        )
+        cv_image = cv2.warpAffine(
+            cv_image,
+            transform_matrix[:2],
+            (self.input_size[1], self.input_size[0]),
+            flags=cv2.INTER_LINEAR,
+        )
+
+        encoder_inputs = {
+            self.encoder_input_name: cv_image.astype(np.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

src/prediction_api/samgeo_predictors.py

@@ -1,56 +0,0 @@
-# Press the green button in the gutter to run the script.
-import json
-import os
-
-from src import app_logger
-from src.utilities.constants import ROOT, MODEL_NAME, ZOOM, SOURCE_TYPE
-from src.utilities.type_hints import input_floatlist
-from src.utilities.utilities import get_system_info
-
-
-def samgeo_fast_predict(
-        bbox: input_floatlist, zoom: float = ZOOM, model_name: str = MODEL_NAME, root_folder: str = ROOT, source_type: str = SOURCE_TYPE, crs="EPSG:4326"
-) -> dict:
-    import tempfile
-    from samgeo import tms_to_geotiff
-    from samgeo.fast_sam import SamGeo
-
-    try:
-        os.environ['MPLCONFIGDIR'] = root_folder
-        get_system_info()
-    except Exception as e:
-        app_logger.error(f"Error while setting 'MPLCONFIGDIR':{e}.")
-
-    with tempfile.NamedTemporaryFile(prefix=f"{source_type}_", suffix=".tif", dir=root_folder) as image_input_tmp:
-        app_logger.info(f"start tms_to_geotiff using bbox:{bbox}, type:{type(bbox)}, download image to {image_input_tmp.name} ...")
-        for coord in bbox:
-            app_logger.info(f"coord:{coord}, type:{type(coord)}.")
-
-        # bbox: image input coordinate
-        tms_to_geotiff(output=image_input_tmp.name, bbox=bbox, zoom=zoom, source=source_type, overwrite=True, crs=crs)
-        app_logger.info(f"geotiff created, start to initialize SamGeo instance (read model {model_name} from {root_folder})...")
-
-        predictor = SamGeo(
-            model_type=model_name,
-            checkpoint_dir=root_folder,
-            automatic=False,
-            sam_kwargs=None,
-        )
-        app_logger.info(f"initialized SamGeo instance, start to use SamGeo.set_image({image_input_tmp.name})...")
-        predictor.set_image(image_input_tmp.name)
-
-        with tempfile.NamedTemporaryFile(prefix="output_", suffix=".tif", dir=root_folder) as image_output_tmp:
-            app_logger.info(f"done set_image, start prediction using {image_output_tmp.name} as output...")
-            predictor.everything_prompt(output=image_output_tmp.name)
-
-            # geotiff to geojson
-            with tempfile.NamedTemporaryFile(prefix="feats_", suffix=".geojson", dir=root_folder) as vector_tmp:
-                app_logger.info(f"done prediction, start conversion SamGeo.tiff_to_geojson({image_output_tmp.name}) => {vector_tmp.name}.")
-                predictor.tiff_to_geojson(image_output_tmp.name, vector_tmp.name, bidx=1)
-
-                app_logger.info(f"start reading geojson {vector_tmp.name}...")
-                with open(vector_tmp.name) as out_gdf:
-                    out_gdf_str = out_gdf.read()
-                    out_gdf_json = json.loads(out_gdf_str)
-                    app_logger.info(f"geojson {vector_tmp.name} string has length: {len(out_gdf_str)}.")
-                    return out_gdf_json

src/utilities/constants.py

@@ -22,7 +22,8 @@ CUSTOM_RESPONSE_MESSAGES = {
     422: "Missing required parameter",
     500: "Internal server error"
 }
-MODEL_NAME = "FastSAM-s.pt"
+MODEL_ENCODER_NAME = "mobile_sam.encoder.onnx"
+MODEL_DECODER_NAME = "sam_vit_h_4b8939.decoder.onnx"
 ZOOM = 13
 SOURCE_TYPE = "Satellite"
 

src/utilities/serialize.py

@@ -0,0 +1,85 @@
+"""Serialize objects"""
+from typing import Mapping
+
+from src import app_logger
+from src.utilities.type_hints import ts_dict_str2, ts_dict_str3
+
+
+def serialize(obj: any, include_none: bool = False) -> object:
+    """
+    Return the input object into a serializable one
+
+    Args:
+        obj: Object to serialize
+        include_none: bool to indicate if include also keys with None values during dict serialization
+
+    Returns:
+        object: serialized object
+
+    """
+    return _serialize(obj, include_none)
+
+
+def _serialize(obj: any, include_none: bool) -> any:
+    import numpy as np
+
+    primitive = (int, float, str, bool)
+    # print(type(obj))
+    try:
+        if obj is None:
+            return None
+        elif isinstance(obj, np.integer):
+            return int(obj)
+        elif isinstance(obj, np.floating):
+            return float(obj)
+        elif isinstance(obj, np.ndarray):
+            return obj.tolist()
+        elif isinstance(obj, primitive):
+            return obj
+        elif type(obj) is list:
+            return _serialize_list(obj, include_none)
+        elif type(obj) is tuple:
+            return list(obj)
+        elif type(obj) is bytes:
+            return _serialize_bytes(obj)
+        elif isinstance(obj, Exception):
+            return _serialize_exception(obj)
+        # elif isinstance(obj, object):
+        #     return _serialize_object(obj, include_none)
+        else:
+            return _serialize_object(obj, include_none)
+    except Exception as e_serialize:
+        app_logger.error(f"e_serialize::{e_serialize}, type_obj:{type(obj)}, obj:{obj}.")
+        return f"object_name:{str(obj)}__object_type_str:{str(type(obj))}."
+
+
+def _serialize_object(obj: Mapping[any, object], include_none: bool) -> dict[any]:
+    from bson import ObjectId
+
+    res = {}
+    if type(obj) is not dict:
+        keys = [i for i in obj.__dict__.keys() if (getattr(obj, i) is not None) or include_none]
+    else:
+        keys = [i for i in obj.keys() if (obj[i] is not None) or include_none]
+    for key in keys:
+        if type(obj) is not dict:
+            res[key] = _serialize(getattr(obj, key), include_none)
+        elif isinstance(obj[key], ObjectId):
+            continue
+        else:
+            res[key] = _serialize(obj[key], include_none)
+    return res
+
+
+def _serialize_list(ls: list, include_none: bool) -> list:
+    return [_serialize(elem, include_none) for elem in ls]
+
+
+def _serialize_bytes(b: bytes) -> ts_dict_str2:
+    import base64
+    encoded = base64.b64encode(b)
+    return {"value": encoded.decode('ascii'), "type": "bytes"}
+
+
+def _serialize_exception(e: Exception) -> ts_dict_str3:
+    return {"msg": str(e), "type": str(type(e)), **e.__dict__}