# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
from edgetpu.detection.engine import DetectionEngine
import numpy as np
from PIL import Image

from . import test_utils

def mobilenet_ssd_v1_coco_engine():
  return  DetectionEngine(
      test_utils.TestDataPath('mobilenet_ssd_v1_coco_quant_postprocess.tflite'))

class TestDetectionEnginePythonAPI(unittest.TestCase):

  def _TestCat(self, model_name):
    engine = DetectionEngine(test_utils.TestDataPath(model_name))
    with test_utils.TestImage('cat.bmp') as img:
      ret = engine.DetectWithImage(img, top_k=1)
      self.assertEqual(len(ret), 1)
      self.assertEqual(ret[0].label_id, 16)  # cat
      self.assertGreater(ret[0].score, 0.7)
      self.assertGreater(
          test_utils.IOU(
              np.array([[0.1, 0.1], [0.7, 1.0]]), ret[0].bounding_box), 0.88)
      # Check coordinates in pixels.
      ret = engine.DetectWithImage(img, top_k=1, relative_coord=False)
      self.assertGreater(
          test_utils.IOU(
              np.array([[60, 40], [420, 400]]), ret[0].bounding_box), 0.88)

  def _TestFace(self, model_name):
    engine = DetectionEngine(test_utils.TestDataPath(model_name))
    with test_utils.TestImage('face.jpg') as img:
      ret = engine.DetectWithImage(img, top_k=1)
      self.assertEqual(len(ret), 1)
      self.assertEqual(ret[0].label_id, 0)
      self.assertGreater(ret[0].score, 0.95)
      self.assertGreater(
          test_utils.IOU(
              np.array([[0.41, 0.07], [0.78, 0.49]]), ret[0].bounding_box), 0.9)

      ret = engine.DetectWithImage(img, top_k=1, relative_coord=False)
      # Check coordinates in pixels.
      self.assertGreater(
          test_utils.IOU(
              np.array([[427, 53], [801, 354]]), ret[0].bounding_box), 0.9)

  def _TestPet(self, model_name):
    engine = DetectionEngine(test_utils.TestDataPath(model_name))
    with test_utils.TestImage('cat.bmp') as img:
      ret = engine.DetectWithImage(img, top_k=1)
      self.assertEqual(len(ret), 1)
      self.assertEqual(ret[0].label_id, 0)  # Abyssinian
      self.assertGreater(ret[0].score, 0.9)
      self.assertGreater(
          test_utils.IOU(
              np.array([[0.35, 11], [0.7, 0.66]]), ret[0].bounding_box), 0.84)

      ret = engine.DetectWithImage(img, top_k=1, relative_coord=False)
      # Check coordinates in pixels.
      self.assertGreater(
          test_utils.IOU(
              np.array([[211, 47], [415, 264]]), ret[0].bounding_box), 0.84)

  def testImageObject(self):
    engine = mobilenet_ssd_v1_coco_engine()
    with test_utils.TestImage('cat.bmp') as img:
      ret = engine.DetectWithImage(img, top_k=1)
      self.assertEqual(len(ret), 1)
      self.assertEqual(ret[0].label_id, 16)  # cat
      self.assertGreater(ret[0].score, 0.79)
      self.assertGreater(
          test_utils.IOU(
              np.array([[0.1, 0.1], [0.7, 1.0]]), ret[0].bounding_box), 0.9)

      # Detect with different resample algorithm.
      ret = engine.DetectWithImage(
          img, top_k=1, resample=Image.BICUBIC)
      self.assertEqual(len(ret), 1)
      self.assertEqual(ret[0].label_id, 16)  # cat
      self.assertGreater(ret[0].score, 0.79)
      self.assertGreater(
          test_utils.IOU(
              np.array([[0.1, 0.1], [0.7, 1.0]]), ret[0].bounding_box), 0.91)

      # No error when top_k > number limit of detection candidates.
      engine.DetectWithImage(img, top_k=100000)

  def testImageObjectWithoutLabels(self):
    engine = mobilenet_ssd_v1_coco_engine()
    with test_utils.TestImage('cat.bmp') as img:
      ret = engine.DetectWithImage(img, top_k=1)
      self.assertEqual(len(ret), 1)
      self.assertEqual(ret[0].label_id, 16)  # cat
      self.assertGreater(ret[0].score, 0.79)
      self.assertGreater(
          test_utils.IOU(
              np.array([[0.1, 0.1], [0.7, 1.0]]), ret[0].bounding_box), 0.9)

  def testRawInput(self):
    engine = mobilenet_ssd_v1_coco_engine()
    with test_utils.TestImage('cat.bmp') as img:
      input_tensor = np.asarray(img.resize((300, 300), Image.NEAREST)).flatten()
      ret = engine.DetectWithInputTensor(input_tensor, top_k=1)
      self.assertEqual(len(ret), 1)
      self.assertEqual(ret[0].label_id, 16)  # cat
      self.assertGreater(ret[0].score, 0.79)
      self.assertGreater(
          test_utils.IOU(
              np.array([[0.1, 0.1], [0.7, 1.0]]), ret[0].bounding_box), 0.9)

  def testVariousModelsWithCat(self):
    for model in ['mobilenet_ssd_v1_coco_quant_postprocess.tflite',
                  'mobilenet_ssd_v1_coco_quant_postprocess_edgetpu.tflite',
                  'mobilenet_ssd_v2_coco_quant_postprocess.tflite',
                  'mobilenet_ssd_v2_coco_quant_postprocess_edgetpu.tflite']:
      self._TestCat(model)

  def testFaceDetection(self):
    for model in ['mobilenet_ssd_v2_face_quant_postprocess.tflite',
                  'mobilenet_ssd_v2_face_quant_postprocess_edgetpu.tflite']:
      self._TestFace(model)

  def testPetDetection(self):
    for model in ['ssd_mobilenet_v1_fine_tuned.tflite',
                  'ssd_mobilenet_v1_fine_tuned_edgetpu.tflite']:
      self._TestPet(model)


if __name__ == '__main__':
  unittest.main()