<commit_before># Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# 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
#
#     http://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.
"""Tests for image_classifier."""

import enum

from absl.testing import parameterized
import tensorflow as tf

from tensorflow_lite_support.python.task.core.proto import base_options_pb2
from tensorflow_lite_support.python.task.processor.proto import bounding_box_pb2
from tensorflow_lite_support.python.task.processor.proto import classification_options_pb2
from tensorflow_lite_support.python.task.processor.proto import classifications_pb2
from tensorflow_lite_support.python.task.vision import image_classifier
from tensorflow_lite_support.python.task.vision.core import tensor_image
from tensorflow_lite_support.python.test import test_util

_BaseOptions = base_options_pb2.BaseOptions
_ImageClassifier = image_classifier.ImageClassifier
_ImageClassifierOptions = image_classifier.ImageClassifierOptions

_MODEL_FILE = 'mobilenet_v2_1.0_224.tflite'
_IMAGE_FILE = 'burger.jpg'
_ALLOW_LIST = ['cheeseburger', 'guacamole']
_DENY_LIST = ['cheeseburger']
_SCORE_THRESHOLD = 0.5
_MAX_RESULTS = 3


def _create_classifier_from_options(base_options, **classification_options):
  classification_options = classification_options_pb2.ClassificationOptions(
      **classification_options)
  options = _ImageClassifierOptions(
      base_options=base_options, classification_options=classification_options)
  classifier = _ImageClassifier.create_from_options(options)
  return classifier


class ModelFileType(enum.Enum):
  FILE_CONTENT = 1
  FILE_NAME = 2


class ImageClassifierTest(parameterized.TestCase, tf.test.TestCase):

  def setUp(self):
    super().setUp()
    self.test_image_path = test_util.get_test_data_path(_IMAGE_FILE)
    self.model_path = test_util.get_test_data_path(_MODEL_FILE)

  def test_create_from_file_succeeds_with_valid_model_path(self):
    # Creates with default option and valid model file successfully.
    classifier = _ImageClassifier.create_from_file(self.model_path)
    self.assertIsInstance(classifier, _ImageClassifier)

  def test_create_from_options_succeeds_with_valid_model_path(self):
    # Creates with options containing model file successfully.
    base_options = _BaseOptions(file_name=self.model_path)
    options = _ImageClassifierOptions(base_options=base_options)
    classifier = _ImageClassifier.create_from_options(options)
    self.assertIsInstance(classifier, _ImageClassifier)

  def test_create_from_options_fails_with_invalid_model_path(self):
    # Invalid empty model path.
    with self.assertRaisesRegex(
        ValueError,
        r"ExternalFile must specify at least one of 'file_content', "
        r"'file_name' or 'file_descriptor_meta'."):
      base_options = _BaseOptions(file_name='')
      options = _ImageClassifierOptions(base_options=base_options)
      _ImageClassifier.create_from_options(options)

  def test_create_from_options_succeeds_with_valid_model_content(self):
    # Creates with options containing model content successfully.
    with open(self.model_path, 'rb') as f:
      base_options = _BaseOptions(file_content=f.read())
      options = _ImageClassifierOptions(base_options=base_options)
      classifier = _ImageClassifier.create_from_options(options)
      self.assertIsInstance(classifier, _ImageClassifier)

  @parameterized.parameters((ModelFileType.FILE_NAME, 3, """
  classifications {
    classes {
      index: 934
      score: 0.739974
      class_name: "cheeseburger"
    }
    classes {
      index: 925
      score: 0.026929
      class_name: "guacamole"
    }
    classes { 
      index: 932 
      score: 0.025737 
      class_name: "bagel" 
    }
    head_index: 0
  }
  """), (ModelFileType.FILE_CONTENT, 3, """
  classifications {
    classes {
      index: 934
      score: 0.739974
      class_name: "cheeseburger"
    }
    classes {
      index: 925
      score: 0.026929
      class_name: "guacamole"
    }
    classes { 
      index: 932 
      score: 0.025737 
      class_name: "bagel" 
    }
    head_index: 0
  }
  """))
  def test_classify_model(self, model_file_type, max_results,
                          expected_result_text_proto):
    # Creates classifier.
    if model_file_type is ModelFileType.FILE_NAME:
      base_options = _BaseOptions(file_name=self.model_path)
    elif model_file_type is ModelFileType.FILE_CONTENT:
      with open(self.model_path, 'rb') as f:
        model_content = f.read()
      base_options = _BaseOptions(file_content=model_content)
    else:
      # Should never happen
      raise ValueError('model_file_type is invalid.')

    classifier = _create_classifier_from_options(
        base_options, max_results=max_results)

    # Loads image.
    image = tensor_image.TensorImage.create_from_file(self.test_image_path)

    # Classifies the input.
    image_result = classifier.classify(image, bounding_box=None)

    # Comparing results (classification w/o bounding box).
    classification_result = classifications_pb2.ClassificationResult()
    classification_result.ParseFromString(image_result.SerializeToString())
    self.assertProtoEquals(expected_result_text_proto, image_result)

  def test_classify_model_with_bounding_box(self):
    # Creates classifier.
    base_options = _BaseOptions(file_name=self.model_path)

    classifier = _create_classifier_from_options(base_options, max_results=3)

    # Loads image.
    image = tensor_image.TensorImage.create_from_file(self.test_image_path)

    # Bounding box in "burger.jpg" corresponding to "burger_crop.jpg".
    bounding_box = bounding_box_pb2.BoundingBox(
        origin_x=0, origin_y=0, width=400, height=325)

    # Classifies the input.
    image_result = classifier.classify(image, bounding_box)

    # Expected results.
    expected_result_text_proto = """
    classifications {
      classes {
        index: 934
        score: 0.881507
        class_name: "cheeseburger"
      }
      classes {
        index: 925
        score: 0.019457
        class_name: "guacamole"
      }
      classes { 
        index: 932 
        score: 0.012489 
        class_name: "bagel" 
      }
      head_index: 0
    }
    """

    # Comparing results (classification w/ bounding box).
    classification_result = classifications_pb2.ClassificationResult()
    classification_result.ParseFromString(image_result.SerializeToString())
    self.assertProtoEquals(expected_result_text_proto, image_result)

  def test_max_results_option(self):
    # Creates classifier.
    base_options = _BaseOptions(file_name=self.model_path)

    classifier = _create_classifier_from_options(
        base_options, max_results=_MAX_RESULTS)

    # Loads image.
    image = tensor_image.TensorImage.create_from_file(self.test_image_path)

    # Classifies the input.
    image_result = classifier.classify(image, bounding_box=None)
    categories = image_result.classifications[0].classes

    self.assertLessEqual(
        len(categories), _MAX_RESULTS, 'Too many results returned.')

  def test_score_threshold_option(self):
    # Creates classifier.
    base_options = _BaseOptions(file_name=self.model_path)

    classifier = _create_classifier_from_options(
        base_options, score_threshold=_SCORE_THRESHOLD)

    # Loads image.<commit_msg>Updated classification with bounding box test<commit_after># Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# 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
#
#     http://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.
"""Tests for image_classifier."""

import enum

from absl.testing import parameterized
import tensorflow as tf

from tensorflow_lite_support.python.task.core.proto import base_options_pb2
from tensorflow_lite_support.python.task.processor.proto import bounding_box_pb2
from tensorflow_lite_support.python.task.processor.proto import classification_options_pb2
from tensorflow_lite_support.python.task.processor.proto import classifications_pb2
from tensorflow_lite_support.python.task.vision import image_classifier
from tensorflow_lite_support.python.task.vision.core import tensor_image
from tensorflow_lite_support.python.test import test_util

_BaseOptions = base_options_pb2.BaseOptions
_ImageClassifier = image_classifier.ImageClassifier
_ImageClassifierOptions = image_classifier.ImageClassifierOptions

_MODEL_FILE = 'mobilenet_v2_1.0_224.tflite'
_IMAGE_FILE = 'burger.jpg'
_ALLOW_LIST = ['cheeseburger', 'guacamole']
_DENY_LIST = ['cheeseburger']
_SCORE_THRESHOLD = 0.5
_MAX_RESULTS = 3


def _create_classifier_from_options(base_options, **classification_options):
  classification_options = classification_options_pb2.ClassificationOptions(
      **classification_options)
  options = _ImageClassifierOptions(
      base_options=base_options, classification_options=classification_options)
  classifier = _ImageClassifier.create_from_options(options)
  return classifier


class ModelFileType(enum.Enum):
  FILE_CONTENT = 1
  FILE_NAME = 2


class ImageClassifierTest(parameterized.TestCase, tf.test.TestCase):

  def setUp(self):
    super().setUp()
    self.test_image_path = test_util.get_test_data_path(_IMAGE_FILE)
    self.model_path = test_util.get_test_data_path(_MODEL_FILE)

  def test_create_from_file_succeeds_with_valid_model_path(self):
    # Creates with default option and valid model file successfully.
    classifier = _ImageClassifier.create_from_file(self.model_path)
    self.assertIsInstance(classifier, _ImageClassifier)

  def test_create_from_options_succeeds_with_valid_model_path(self):
    # Creates with options containing model file successfully.
    base_options = _BaseOptions(file_name=self.model_path)
    options = _ImageClassifierOptions(base_options=base_options)
    classifier = _ImageClassifier.create_from_options(options)
    self.assertIsInstance(classifier, _ImageClassifier)

  def test_create_from_options_fails_with_invalid_model_path(self):
    # Invalid empty model path.
    with self.assertRaisesRegex(
        ValueError,
        r"ExternalFile must specify at least one of 'file_content', "
        r"'file_name' or 'file_descriptor_meta'."):
      base_options = _BaseOptions(file_name='')
      options = _ImageClassifierOptions(base_options=base_options)
      _ImageClassifier.create_from_options(options)

  def test_create_from_options_succeeds_with_valid_model_content(self):
    # Creates with options containing model content successfully.
    with open(self.model_path, 'rb') as f:
      base_options = _BaseOptions(file_content=f.read())
      options = _ImageClassifierOptions(base_options=base_options)
      classifier = _ImageClassifier.create_from_options(options)
      self.assertIsInstance(classifier, _ImageClassifier)

  @parameterized.parameters((ModelFileType.FILE_NAME, 3, """
  classifications {
    classes {
      index: 934
      score: 0.739974
      class_name: "cheeseburger"
    }
    classes {
      index: 925
      score: 0.026929
      class_name: "guacamole"
    }
    classes { 
      index: 932 
      score: 0.025737 
      class_name: "bagel" 
    }
    head_index: 0
  }
  """), (ModelFileType.FILE_CONTENT, 3, """
  classifications {
    classes {
      index: 934
      score: 0.739974
      class_name: "cheeseburger"
    }
    classes {
      index: 925
      score: 0.026929
      class_name: "guacamole"
    }
    classes { 
      index: 932 
      score: 0.025737 
      class_name: "bagel" 
    }
    head_index: 0
  }
  """))
  def test_classify_model(self, model_file_type, max_results,
                          expected_result_text_proto):
    # Creates classifier.
    if model_file_type is ModelFileType.FILE_NAME:
      base_options = _BaseOptions(file_name=self.model_path)
    elif model_file_type is ModelFileType.FILE_CONTENT:
      with open(self.model_path, 'rb') as f:
        model_content = f.read()
      base_options = _BaseOptions(file_content=model_content)
    else:
      # Should never happen
      raise ValueError('model_file_type is invalid.')

    classifier = _create_classifier_from_options(
        base_options, max_results=max_results)

    # Loads image.
    image = tensor_image.TensorImage.create_from_file(self.test_image_path)

    # Classifies the input.
    image_result = classifier.classify(image, bounding_box=None)

    # Comparing results (classification w/o bounding box).
    classification_result = classifications_pb2.ClassificationResult()
    classification_result.ParseFromString(image_result.SerializeToString())
    self.assertProtoEquals(expected_result_text_proto, image_result)

  def test_classify_model_with_bounding_box(self):
    # Creates classifier.
    base_options = _BaseOptions(file_name=self.model_path)

    classifier = _create_classifier_from_options(base_options, max_results=3)

    # Loads image.
    image = tensor_image.TensorImage.create_from_file(self.test_image_path)

    # Bounding box in "burger.jpg" corresponding to "burger_crop.jpg".
    bounding_box = bounding_box_pb2.BoundingBox(
        origin_x=0, origin_y=0, width=400, height=325)

    # Classifies the input.
    image_result = classifier.classify(image, bounding_box)

    # Expected results.
    expected_result_text_proto = """
    classifications {
      classes {
        index: 934
        score: 0.881507
        class_name: "cheeseburger"
      }
      classes {
        index: 925
        score: 0.019457
        class_name: "guacamole"
      }
      classes { 
        index: 932 
        score: 0.012489 
        class_name: "bagel" 
      }
      head_index: 0
    }
    """

    # Comparing results (classification w/ bounding box).
    self.assertProtoEquals(expected_result_text_proto, image_result)

  def test_max_results_option(self):
    # Creates classifier.
    base_options = _BaseOptions(file_name=self.model_path)

    classifier = _create_classifier_from_options(
        base_options, max_results=_MAX_RESULTS)

    # Loads image.
    image = tensor_image.TensorImage.create_from_file(self.test_image_path)

    # Classifies the input.
    image_result = classifier.classify(image, bounding_box=None)
    categories = image_result.classifications[0].classes

    self.assertLessEqual(
        len(categories), _MAX_RESULTS, 'Too many results returned.')

  def test_score_threshold_option(self):
    # Creates classifier.
    base_options = _BaseOptions(file_name=self.model_path)

    classifier = _create_classifier_from_options(
        base_options, score_threshold=_SCORE_THRESHOLD)

    # Loads image.
{'fraction_non_alphanumeric': 0.07558177221754067, 'fraction_numerical': 0.017177675503986516, 'mean_word_length': 4.093430174975239, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3438186', 'n_tokens_mistral': 4951, 'n_tokens_neox': 4620, 'n_words': 957}