# RUN: %PYTHON %s
# XFAIL: *

import absl
import absl.testing
import iree.compiler.tflite as iree_tflite_compile
import iree.runtime as iree_rt
import numpy
import test_util
import urllib.request

from PIL import Image

model_path = "https://github.com/tensorflow/tflite-micro/raw/main/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=2).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/main/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])))
    args = [im.reshape(shape)]
    return args

  def test_compile_tflite(self):
    # tflite.interpreter python API has problem rendering this file. Issue filed.
    # The example would fail after the iree_tflite_compile.compile_file.
    self.compile_and_execute()

if __name__ == '__main__':
  absl.testing.absltest.main()

# RUN: %PYTHON %s
# XFAIL: *

import absl.testing
import numpy
import test_util
import urllib.request

from PIL import Image

model_path = "https://github.com/tensorflow/tflite-micro/raw/main/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/main/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])))
    args = [im.reshape(shape)]
    return args

  def test_compile_tflite(self):
    # tflite.interpreter python API has problem rendering this file. Issue filed.
    # The example would fail after the iree_tflite_compile.compile_file.
    self.compile_and_execute()

if __name__ == '__main__':
  absl.testing.absltest.main()

Remove unused import modules in person_detect_test