# Path: pycoral/adapters/classify.py
# def num_classes(interpreter):
# def get_scores(interpreter):
# def get_classes_from_scores(scores,
#                             top_k=float('inf'),
#                             score_threshold=-float('inf')):
# def get_classes(interpreter, top_k=float('inf'), score_threshold=-float('inf')):
# 
# Path: pycoral/adapters/common.py
# def output_tensor(interpreter, i):
# def input_details(interpreter, key):
# def input_size(interpreter):
# def input_tensor(interpreter):
# def set_input(interpreter, data):
# def set_resized_input(interpreter, size, resize):
# 
# Path: pycoral/learn/imprinting/engine.py
# class ImprintingEngine:
# 
#   def __init__(self, model_path, keep_classes=False):
#     """Performs weight imprinting (transfer learning) with the given model.
# 
#     Args:
#       model_path (str): Path to the ``.tflite`` model you want to retrain.
#         This must be a model that's specially-designed for this API. You
#         can use our `weight imprinting model
#         <https://coral.ai/models/image-classification/#training-models>`_ that
#         has a pre-trained base model, or you can train the base model yourself
#         by following our guide to `Retrain the base MobileNet model
#         <https://coral.ai/docs/edgetpu/retrain-classification-ondevice/#retrain-the-base-mobilenet-model>`_.
#       keep_classes (bool): If True, keep the existing classes from the
#         pre-trained model (and use training to add additional classes). If
#         False, drop the existing classes and train the model to include new
#         classes only.
#     """
#     self._engine = _pywrap_coral.ImprintingEnginePythonWrapper(
#         model_path, keep_classes)
# 
#   @property
#   def embedding_dim(self):
#     """Returns number of embedding dimensions."""
#     return self._engine.EmbeddingDim()
# 
#   @property
#   def num_classes(self):
#     """Returns number of currently trained classes."""
#     return self._engine.NumClasses()
# 
#   def serialize_extractor_model(self):
#     """Returns embedding extractor model as `bytes` object."""
#     return self._engine.SerializeExtractorModel()
# 
#   def serialize_model(self):
#     """Returns newly trained model as `bytes` object."""
#     return self._engine.SerializeModel()
# 
#   def train(self, embedding, class_id):
#     """Trains the model with the given embedding for specified class.
# 
#     You can use this to add new classes to the model or retrain classes that you
#     previously added using this imprinting API.
# 
#     Args:
#       embedding (:obj:`numpy.array`): The embedding vector for training
#         specified single class.
#       class_id (int): The label id for this class. The index must be either the
#         number of existing classes (to add a new class to the model) or the
#         index of an existing class that was trained using this imprinting API
#         (you can't retrain classes from the pre-trained model).
#     """
#     self._engine.Train(embedding, class_id)
# 
# Path: pycoral/utils/edgetpu.py
# def make_interpreter(model_path_or_content, device=None, delegate=None):
#   """Creates a new ``tf.lite.Interpreter`` instance using the given model.
# 
#   **Note:** If you have multiple Edge TPUs, you should always specify the
#   ``device`` argument.
# 
#   Args:
#      model_path_or_content (str or bytes): `str` object is interpreted as
#        model path, `bytes` object is interpreted as model content.
#      device (str): The Edge TPU device you want:
# 
#        + None      -- use any Edge TPU (this is the default)
#        + ":<N>"    -- use N-th Edge TPU (this corresponds to the enumerated
#          index position from :func:`list_edge_tpus`)
#        + "usb"     -- use any USB Edge TPU
#        + "usb:<N>" -- use N-th USB Edge TPU
#        + "pci"     -- use any PCIe Edge TPU
#        + "pci:<N>" -- use N-th PCIe Edge TPU
# 
#        If left as None, you cannot reliably predict which device you'll get.
#        So if you have multiple Edge TPUs and want to run a specific model on
#        each one, then you must specify the device.
#      delegate: A pre-loaded Edge TPU delegate object, as provided by
#        :func:`load_edgetpu_delegate`. If provided, the `device` argument
#        is ignored.
# 
#   Returns:
#      New ``tf.lite.Interpreter`` instance.
#   """
#   if delegate:
#     delegates = [delegate]
#   else:
#     delegates = [load_edgetpu_delegate({'device': device} if device else {})]
#   if isinstance(model_path_or_content, bytes):
#     return tflite.Interpreter(
#         model_content=model_path_or_content, experimental_delegates=delegates)
#   else:
#     return tflite.Interpreter(
#         model_path=model_path_or_content, experimental_delegates=delegates)


import argparse
import os
import numpy as np
from PIL import Image
from pycoral.adapters import classify
from pycoral.adapters import common
from pycoral.learn.imprinting.engine import ImprintingEngine
from pycoral.utils.edgetpu import make_interpreter
  print('----------------      Start training     -----------------')
  num_classes = engine.num_classes
  for class_id, tensors in enumerate(train_input):
    for tensor in tensors:
      common.set_input(extractor, tensor)
      extractor.invoke()
      embedding = classify.get_scores(extractor)
      engine.train(embedding, class_id=num_classes + class_id)
  print('----------------     Training finished!  -----------------')

  with open(args.output, 'wb') as f:
    f.write(engine.serialize_model())
  print('Model saved as : ', args.output)
  _save_labels(labels_map, args.output)

  print('------------------   Start evaluating   ------------------')
  interpreter = make_interpreter(args.output)
  interpreter.allocate_tensors()
  size = common.input_size(interpreter)

  top_k = 5
  correct = [0] * top_k
  wrong = [0] * top_k
  for category, image_list in test_set.items():
    print('Evaluating category [', category, ']')
    for img_name in image_list:
      img = Image.open(os.path.join(args.data, category,
                                    img_name)).resize(size, Image.NEAREST)
      common.set_input(interpreter, img)
      interpreter.invoke()

# Path: pycoral/adapters/classify.py
# def num_classes(interpreter):
# def get_scores(interpreter):
# def get_classes_from_scores(scores,
#                             top_k=float('inf'),
#                             score_threshold=-float('inf')):
# def get_classes(interpreter, top_k=float('inf'), score_threshold=-float('inf')):
# 
# Path: pycoral/adapters/common.py
# def output_tensor(interpreter, i):
# def input_details(interpreter, key):
# def input_size(interpreter):
# def input_tensor(interpreter):
# def set_input(interpreter, data):
# def set_resized_input(interpreter, size, resize):
# 
# Path: pycoral/learn/imprinting/engine.py
# class ImprintingEngine:
# 
#   def __init__(self, model_path, keep_classes=False):
#     """Performs weight imprinting (transfer learning) with the given model.
# 
#     Args:
#       model_path (str): Path to the ``.tflite`` model you want to retrain.
#         This must be a model that's specially-designed for this API. You
#         can use our `weight imprinting model
#         <https://coral.ai/models/image-classification/#training-models>`_ that
#         has a pre-trained base model, or you can train the base model yourself
#         by following our guide to `Retrain the base MobileNet model
#         <https://coral.ai/docs/edgetpu/retrain-classification-ondevice/#retrain-the-base-mobilenet-model>`_.
#       keep_classes (bool): If True, keep the existing classes from the
#         pre-trained model (and use training to add additional classes). If
#         False, drop the existing classes and train the model to include new
#         classes only.
#     """
#     self._engine = _pywrap_coral.ImprintingEnginePythonWrapper(
#         model_path, keep_classes)
# 
#   @property
#   def embedding_dim(self):
#     """Returns number of embedding dimensions."""
#     return self._engine.EmbeddingDim()
# 
#   @property
#   def num_classes(self):
#     """Returns number of currently trained classes."""
#     return self._engine.NumClasses()
# 
#   def serialize_extractor_model(self):
#     """Returns embedding extractor model as `bytes` object."""
#     return self._engine.SerializeExtractorModel()
# 
#   def serialize_model(self):
#     """Returns newly trained model as `bytes` object."""
#     return self._engine.SerializeModel()
# 
#   def train(self, embedding, class_id):
#     """Trains the model with the given embedding for specified class.
# 
#     You can use this to add new classes to the model or retrain classes that you
#     previously added using this imprinting API.
# 
#     Args:
#       embedding (:obj:`numpy.array`): The embedding vector for training
#         specified single class.
#       class_id (int): The label id for this class. The index must be either the
#         number of existing classes (to add a new class to the model) or the
#         index of an existing class that was trained using this imprinting API
#         (you can't retrain classes from the pre-trained model).
#     """
#     self._engine.Train(embedding, class_id)
# 
# Path: pycoral/utils/edgetpu.py
# def make_interpreter(model_path_or_content, device=None, delegate=None):
#   """Creates a new ``tf.lite.Interpreter`` instance using the given model.
# 
#   **Note:** If you have multiple Edge TPUs, you should always specify the
#   ``device`` argument.
# 
#   Args:
#      model_path_or_content (str or bytes): `str` object is interpreted as
#        model path, `bytes` object is interpreted as model content.
#      device (str): The Edge TPU device you want:
# 
#        + None      -- use any Edge TPU (this is the default)
#        + ":<N>"    -- use N-th Edge TPU (this corresponds to the enumerated
#          index position from :func:`list_edge_tpus`)
#        + "usb"     -- use any USB Edge TPU
#        + "usb:<N>" -- use N-th USB Edge TPU
#        + "pci"     -- use any PCIe Edge TPU
#        + "pci:<N>" -- use N-th PCIe Edge TPU
# 
#        If left as None, you cannot reliably predict which device you'll get.
#        So if you have multiple Edge TPUs and want to run a specific model on
#        each one, then you must specify the device.
#      delegate: A pre-loaded Edge TPU delegate object, as provided by
#        :func:`load_edgetpu_delegate`. If provided, the `device` argument
#        is ignored.
# 
#   Returns:
#      New ``tf.lite.Interpreter`` instance.
#   """
#   if delegate:
#     delegates = [delegate]
#   else:
#     delegates = [load_edgetpu_delegate({'device': device} if device else {})]
#   if isinstance(model_path_or_content, bytes):
#     return tflite.Interpreter(
#         model_content=model_path_or_content, experimental_delegates=delegates)
#   else:
#     return tflite.Interpreter(
#         model_path=model_path_or_content, experimental_delegates=delegates)


# Path: examples/imprinting_learning.py
import argparse
import os
import numpy as np
from PIL import Image
from pycoral.adapters import classify
from pycoral.adapters import common
from pycoral.learn.imprinting.engine import ImprintingEngine
from pycoral.utils.edgetpu import make_interpreter

  print('----------------      Start training     -----------------')
  num_classes = engine.num_classes
  for class_id, tensors in enumerate(train_input):
    for tensor in tensors:
      common.set_input(extractor, tensor)
      extractor.invoke()
      embedding = classify.get_scores(extractor)
      engine.train(embedding, class_id=num_classes + class_id)
  print('----------------     Training finished!  -----------------')

  with open(args.output, 'wb') as f:
    f.write(engine.serialize_model())
  print('Model saved as : ', args.output)
  _save_labels(labels_map, args.output)

  print('------------------   Start evaluating   ------------------')
  interpreter = make_interpreter(args.output)
  interpreter.allocate_tensors()
  size = common.input_size(interpreter)

  top_k = 5
  correct = [0] * top_k
  wrong = [0] * top_k
  for category, image_list in test_set.items():
    print('Evaluating category [', category, ']')
    for img_name in image_list:
      img = Image.open(os.path.join(args.data, category,
                                    img_name)).resize(size, Image.NEAREST)
      common.set_input(interpreter, img)
      interpreter.invoke()

      candidates = classify.get_classes(interpreter, top_k, score_threshold=0.1)