File size: 3,251 Bytes
94729d4 5569a41 94729d4 c8c2cad 94729d4 5569a41 94729d4 5569a41 94729d4 5569a41 94729d4 5569a41 84a883f 5569a41 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import json
import unittest
from unittest.mock import patch
import numpy as np
from samgis.prediction_api import predictors
from samgis.prediction_api.predictors import get_raster_inference, samexporter_predict
from tests import TEST_EVENTS_FOLDER
class TestPredictors(unittest.TestCase):
@patch.object(predictors, "SegmentAnythingONNX")
def test_get_raster_inference(self, segment_anything_onnx_mocked):
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():
model_mocked = segment_anything_onnx_mocked()
img = np.load(TEST_EVENTS_FOLDER / f"{name_fn}" / k / "img.npy")
inference_out = np.load(TEST_EVENTS_FOLDER / f"{name_fn}" / k / "inference_out.npy")
mask = np.load(TEST_EVENTS_FOLDER / f"{name_fn}" / k / "mask.npy")
prompt = input_output["input"]["prompt"]
model_name = input_output["input"]["model_name"]
model_mocked.embed.return_value = np.array(img)
model_mocked.embed.side_effect = None
model_mocked.predict_masks.return_value = inference_out
model_mocked.predict_masks.side_effect = None
print(f"k:{k}.")
output_mask, len_inference_out = get_raster_inference(
img=img,
prompt=prompt,
models_instance=model_mocked,
model_name=model_name
)
assert np.array_equal(output_mask, mask)
assert len_inference_out == input_output["output"]["n_predictions"]
@patch.object(predictors, "get_raster_inference_with_embedding_from_dict")
@patch.object(predictors, "SegmentAnythingONNX")
@patch.object(predictors, "download_extent")
@patch.object(predictors, "get_vectorized_raster_as_geojson")
def test_samexporter_predict(
self,
get_vectorized_raster_as_geojson_mocked,
download_extent_mocked,
segment_anything_onnx_mocked,
get_raster_inference_with_embedding_from_dict_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_raster_inference_with_embedding_from_dict_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", source_name="localtest"
)
assert output == {"n_predictions": 1, "geojson": "{}", "n_shapes_geojson": 2}
|