[test] add test cases for get_vectorized_raster_as_geojson and SegmentAnythingONNX

src/io/geo_helpers.py

@@ -31,7 +31,7 @@ def load_affine_transformation_from_matrix(matrix_source_coeffs: List[float]) ->
         raise e
 
 
-def get_affine_transform_from_gdal(matrix_source_coeffs: List[float]) -> Affine:
+def get_affine_transform_from_gdal(matrix_source_coeffs: List[float] or Tuple[float]) -> Affine:
     """wrapper for rasterio Affine from_gdal method
 
     Args:

src/prediction_api/predictors.py

@@ -1,4 +1,4 @@
-"""functions using the machine learning instance model"""
+"""functions using machine learning instance model(s)"""
 from typing import Dict, Tuple
 from PIL.Image import Image
 import numpy as np
@@ -10,6 +10,7 @@ from src.prediction_api.sam_onnx import SegmentAnythingONNX
 from src.utilities.constants import MODEL_ENCODER_NAME, MODEL_DECODER_NAME, DEFAULT_TMS
 from src.utilities.type_hints import llist_float
 
+
 models_dict = {"fastsam": {"instance": None}}
 
 
@@ -36,34 +37,31 @@ def samexporter_predict(
     Returns:
         dict: Affine transform
     """
-    try:
-        if models_dict[model_name]["instance"] is None:
-            app_logger.info(f"missing instance model {model_name}, instantiating it now!")
-            model_instance = SegmentAnythingONNX(
-                encoder_model_path=MODEL_FOLDER / MODEL_ENCODER_NAME,
-                decoder_model_path=MODEL_FOLDER / MODEL_DECODER_NAME
-            )
-            models_dict[model_name]["instance"] = model_instance
-        app_logger.debug(f"using a {model_name} instance model...")
-        models_instance = models_dict[model_name]["instance"]
+    if models_dict[model_name]["instance"] is None:
+        app_logger.info(f"missing instance model {model_name}, instantiating it now!")
+        model_instance = SegmentAnythingONNX(
+            encoder_model_path=MODEL_FOLDER / MODEL_ENCODER_NAME,
+            decoder_model_path=MODEL_FOLDER / MODEL_DECODER_NAME
+        )
+        models_dict[model_name]["instance"] = model_instance
+    app_logger.debug(f"using a {model_name} instance model...")
+    models_instance = models_dict[model_name]["instance"]
 
-        app_logger.info(f'tile_source: {DEFAULT_TMS}!')
-        pt0, pt1 = bbox
-        app_logger.info(f"downloading geo-referenced raster with bbox {bbox}, zoom {zoom}.")
-        img, matrix = download_extent(DEFAULT_TMS, pt0[0], pt0[1], pt1[0], pt1[1], zoom)
-        app_logger.info(f"img type {type(img)} with shape/size:{img.size}, matrix:{type(matrix)}, matrix:{matrix}.")
+    app_logger.info(f'tile_source: {DEFAULT_TMS}!')
+    pt0, pt1 = bbox
+    app_logger.info(f"downloading geo-referenced raster with bbox {bbox}, zoom {zoom}.")
+    img, matrix = download_extent(DEFAULT_TMS, pt0[0], pt0[1], pt1[0], pt1[1], zoom)
+    app_logger.info(f"img type {type(img)} with shape/size:{img.size}, matrix:{type(matrix)}, matrix:{matrix}.")
 
-        transform = get_affine_transform_from_gdal(matrix)
-        app_logger.debug(f"transform to consume with rasterio.shapes: {type(transform)}, {transform}.")
+    transform = get_affine_transform_from_gdal(matrix)
+    app_logger.debug(f"transform to consume with rasterio.shapes: {type(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...")
-        return {
-            "n_predictions": n_predictions,
-            **get_vectorized_raster_as_geojson(mask, matrix)
-        }
-    except ImportError as e_import_module:
-        app_logger.error(f"Error trying import module:{e_import_module}.")
+    mask, n_predictions = get_raster_inference(img, prompt, models_instance, model_name)
+    app_logger.info(f"created {n_predictions} masks, preparing conversion to geojson...")
+    return {
+        "n_predictions": n_predictions,
+        **get_vectorized_raster_as_geojson(mask, matrix)
+    }
 
 
 def get_raster_inference(

tests/io/test_geo_helpers.py

@@ -62,7 +62,7 @@ class TestGeoHelpers(unittest.TestCase):
                                              "check https://github.com/rasterio/affine project for updates")
                     raise e
 
-    def test_get_vectorized_raster_as_geojson(self):
+    def test_get_vectorized_raster_as_geojson_ok(self):
         from src.io.geo_helpers import get_vectorized_raster_as_geojson
 
         name_fn = "samexporter_predict"
@@ -78,3 +78,26 @@ class TestGeoHelpers(unittest.TestCase):
                 output_geojson = shapely.from_geojson(output["geojson"])
                 expected_output_geojson = shapely.from_geojson(input_output["output"]["geojson"])
                 assert shapely.equals_exact(output_geojson, expected_output_geojson, tolerance=0.000006)
+
+    def test_get_vectorized_raster_as_geojson_fail(self):
+        from src.io.geo_helpers import get_vectorized_raster_as_geojson
+
+        name_fn = "samexporter_predict"
+
+        with open(TEST_EVENTS_FOLDER / f"{name_fn}.json") as tst_json:
+            inputs_outputs = json.load(tst_json)
+            for k, input_output in inputs_outputs.items():
+                print(f"k:{k}.")
+                mask = np.load(TEST_EVENTS_FOLDER / name_fn / k / "mask.npy")
+
+                # Could be also another generic Exception, here we intercept TypeError caused by wrong matrix input on
+                # rasterio.Affine.from_gdal() wrapped by get_affine_transform_from_gdal()
+                with self.assertRaises(TypeError):
+                    try:
+                        wrong_matrix = 1.0,
+                        get_vectorized_raster_as_geojson(mask=mask, matrix=wrong_matrix)
+                    except TypeError as te:
+                        print(f"te:{te}.")
+                        msg = "Affine.from_gdal() missing 5 required positional arguments: 'a', 'b', 'f', 'd', and 'e'"
+                        self.assertEqual(str(te), msg)
+                        raise te

tests/prediction_api/test_predictors.py

@@ -2,16 +2,15 @@ import json
 from unittest.mock import patch
 
 import numpy as np
+import rasterio
 
-from src.prediction_api import sam_onnx
-from src.prediction_api.predictors import get_raster_inference
+from src.prediction_api import predictors
+from src.prediction_api.predictors import get_raster_inference, samexporter_predict
 from tests import TEST_EVENTS_FOLDER
 
 
-@patch.object(sam_onnx, "SegmentAnythingONNX")
-def test_get_raster_inference(
-    segment_anything_onnx_mocked
-):
+@patch.object(predictors, "SegmentAnythingONNX")
+def test_get_raster_inference(segment_anything_onnx_mocked):
     name_fn = "samexporter_predict"
 
     with open(TEST_EVENTS_FOLDER / f"{name_fn}.json") as tst_json:
@@ -38,3 +37,32 @@ def test_get_raster_inference(
             )
             assert np.array_equal(output_mask, mask)
             assert len_inference_out == input_output["output"]["n_predictions"]
+
+
+@patch.object(predictors, "get_raster_inference")
+@patch.object(predictors, "SegmentAnythingONNX")
+@patch.object(predictors, "download_extent")
+@patch.object(predictors, "get_vectorized_raster_as_geojson")
+@patch.object(predictors, "get_affine_transform_from_gdal")
+def test_samexporter_predict(
+        get_affine_transform_from_gdal_mocked,
+        get_vectorized_raster_as_geojson_mocked,
+        download_extent_mocked,
+        segment_anything_onnx_mocked,
+        get_raster_inference_mocked
+):
+    """
+    model_instance = SegmentAnythingONNX()
+    img, matrix = download_extent(DEFAULT_TMS, pt0[0], pt0[1], pt1[0], pt1[1], zoom)
+    transform = get_affine_transform_from_gdal(matrix)
+    mask, n_predictions = get_raster_inference(img, prompt, models_instance, model_name)
+    get_vectorized_raster_as_geojson(mask, matrix)
+    """
+    aff = 1, 2, 3, 4, 5, 6
+    segment_anything_onnx_mocked.return_value = "SegmentAnythingONNX_instance"
+    download_extent_mocked.return_value = np.zeros((10, 10)), aff
+    get_affine_transform_from_gdal_mocked.return_value = rasterio.Affine.from_gdal(*aff)
+    get_raster_inference_mocked.return_value = np.ones((10, 10)), 1
+    get_vectorized_raster_as_geojson_mocked.return_value = {"geojson": "{}", "n_shapes_geojson": 2}
+    output = samexporter_predict(bbox=[[1, 2], [3, 4]], prompt=[{}], zoom=10, model_name="fastsam")
+    assert output == {"n_predictions": 1, "geojson": "{}", "n_shapes_geojson": 2}