Files
AMD-SHARK-Studio/tank/tflite/east_text_detector_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

38 lines
1.3 KiB
Python

# RUN: %PYTHON %s
import absl.testing
import numpy
import test_util
model_path = "https://tfhub.dev/sayakpaul/lite-model/east-text-detector/dr/1?lite-format=tflite"
class EastTextDetectorTest(test_util.TFLiteModelTest):
def __init__(self, *args, **kwargs):
super(EastTextDetectorTest, self).__init__(model_path, *args, **kwargs)
def compare_results(self, iree_results, tflite_results, details):
super(EastTextDetectorTest, self).compare_results(
iree_results, tflite_results, details
)
self.assertTrue(
numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all()
)
# The second return is extremely noisy as it is not a binary classification. To handle we
# check normalized correlation with an expectation of "close enough".
iree_norm = numpy.sqrt(iree_results[1] * iree_results[1])
tflite_norm = numpy.sqrt(tflite_results[1] * tflite_results[1])
correlation = numpy.average(
iree_results[1] * tflite_results[1] / iree_norm / tflite_norm
)
self.assertTrue(numpy.isclose(correlation, 1.0, atol=1e-2).all())
def test_compile_tflite(self):
self.compile_and_execute()
if __name__ == "__main__":
absl.testing.absltest.main()