File size: 1,571 Bytes
94729d4 3041467 94729d4 6c70a52 94729d4 6c70a52 |
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 |
import json
from unittest.mock import patch
import numpy as np
from src.prediction_api import sam_onnx
from src.prediction_api.predictors import get_raster_inference
from tests import TEST_EVENTS_FOLDER
@patch.object(sam_onnx, "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:
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"]
|