|
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"] |
|
|
