Files
AMD-SHARK-Studio/tank/tflite/person_detect_test.py
Prashant Kumar 0dcf387089 Add shark_importer for torch_models. (#183)
All the torch_models are imported to gs::shark_tank.
Scripts have been updated.
2022-07-12 20:38:19 -07:00

63 lines
2.0 KiB
Python

# RUN: %PYTHON %s
import absl.testing
import numpy
import test_util
import urllib.request
from PIL import Image
model_path = "https://github.com/tensorflow/tflite-micro/raw/aeac6f39e5c7475cea20c54e86d41e3a38312546/tensorflow/lite/micro/models/person_detect.tflite"
class PersonDetectTest(test_util.TFLiteModelTest):
def __init__(self, *args, **kwargs):
super(PersonDetectTest, self).__init__(model_path, *args, **kwargs)
def compare_results(self, iree_results, tflite_results, details):
super(PersonDetectTest, self).compare_results(
iree_results, tflite_results, details
)
self.assertTrue(
numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all()
)
# TFLite is broken with this model so we hardcode the input/output details.
def setup_tflite(self):
self.input_details = [
{
"shape": [1, 96, 96, 1],
"dtype": numpy.int8,
"index": 0,
}
]
self.output_details = [
{
"shape": [1, 2],
"dtype": numpy.int8,
}
]
# The input has known expected values. We hardcode this value.
def invoke_tflite(self, args):
return [numpy.array([[-113, 113]], dtype=numpy.int8)]
def generate_inputs(self, input_details):
img_path = "https://github.com/tensorflow/tflite-micro/raw/aeac6f39e5c7475cea20c54e86d41e3a38312546/tensorflow/lite/micro/examples/person_detection/testdata/person.bmp"
local_path = "/".join([self.workdir, "person.bmp"])
urllib.request.urlretrieve(img_path, local_path)
shape = input_details[0]["shape"]
im = numpy.array(
Image.open(local_path).resize((shape[1], shape[2]))
).astype(input_details[0]["dtype"])
args = [im.reshape(shape)]
return args
def test_compile_tflite(self):
self.compile_and_execute()
if __name__ == "__main__":
absl.testing.absltest.main()