edgetpu.classification.engine

An inference engine that performs image classification.

class edgetpu.classification.engine.ClassificationEngine(model_path, device_path=None)

Extends BasicEngine to perform image classification with a given model.

This API assumes the given model is trained for image classification.

Parameters:
  • model_path (str) – Path to a TensorFlow Lite (.tflite) file. This model must be compiled for the Edge TPU; otherwise, it simply executes on the host CPU.
  • device_path (str) – The device path for the Edge TPU this engine should use. This argument is needed only when you have multiple Edge TPUs and more inference engines than available Edge TPUs. For details, read how to use multiple Edge TPUs.
Raises:

ValueError – If the model’s output tensor size is not 1.

ClassifyWithImage(img, threshold=0.1, top_k=3, resample=0)

Performs classification with an image.

Parameters:
  • img (PIL.Image) – The image you want to classify.
  • threshold (float) – Minimum confidence threshold for returned classifications. For example, use 0.5 to receive only classifications with a confidence equal-to or higher-than 0.5.
  • top_k (int) – The maximum number of classifications to return.
  • resample (int) – A resampling filter for image resizing. This can be one of PIL.Image.NEAREST, PIL.Image.BOX, PIL.Image.BILINEAR, PIL.Image.HAMMING, PIL.Image.BICUBIC, or PIL.Image.LANCZOS. Default is PIL.Image.NEAREST. See Pillow filters. (Note: A complex filter such as PIL.Image.BICUBIC may create slightly better accuracy but it also causes higher latency.)
Returns:

A list of classifications, each of which is a list [int, float] that represents the label id (int) and the confidence score (float).

Raises:
  • RuntimeError – If the model’s input tensor shape doesn’t match the shape expected for an image classification model, which is [1, height, width, 3].
  • ValueError – If argument values are invalid.
ClassifyWithInputTensor(input_tensor, threshold=0.0, top_k=3)

Performs classification with a raw input tensor.

This requires you to process the input data (the image) and convert it to the appropriately formatted input tensor for your model.

Parameters:
  • input_tensor (numpy.ndarray) – A 1-D array as the input tensor.
  • threshold (float) – Minimum confidence threshold for returned classifications. For example, use 0.5 to receive only classifications with a confidence equal-to or higher-than 0.5.
  • top_k (int) – The maximum number of classifications to return.
Returns:

A list of classifications, each of which is a list [int, float] that represents the label id (int) and the confidence score (float).

Raises:

ValueError – If argument values are invalid.

RunInference(input)

Performs inference with a raw input tensor.

Parameters:
  • input_tensor (numpy.ndarray) – A 1-D array as the input tensor. You can query the required size for this array with required_input_array_size().
  • threshold (float) – Minimum confidence threshold for returned inferences. For example, use 0.5 to receive only inferences with a confidence equal-to or higher-than 0.5.
  • top_k (int) – The maximum number of inferences to return.
Returns:

A 2-tuple with the inference latency in milliseconds (float) and a 1-D array (numpy.ndarray) representing the output tensor. If there are multiple output tensors, they are compressed into a single 1-D array. For example, if the model outputs 2 tensors with values [1, 2, 3] and [0.1, 0.4, 0.9], the returned 1-D array is [1, 2, 3, 0.1, 0.4, 0.9]. You can calculate the size and offset for each tensor using get_all_output_tensors_sizes(), get_num_of_output_tensors(), and get_output_tensor_size().

device_path()

Gets the path for the Edge TPU that’s associated with this inference engine.

See how to run multiple models with multiple Edge TPUs.

Returns:A string representing this engine’s Edge TPU device path.
get_all_output_tensors_sizes()

Gets the size of each output tensor.

A model may output several tensors, but the return from RunInference() and get_raw_output() concatenates them together into a 1-D array. So this function provides the size for each original output tensor, allowing you to calculate the offset for each tensor within the concatenated array.

Returns:An array (numpy.ndarray) with the length of each output tensor (this assumes that all output tensors are 1-D).
get_inference_time()

Gets the latency of the most recent inference.

This can be used by higher level engines for debugging.

Returns:A float representing the inference latency (in milliseconds).
get_input_tensor_shape()

Gets the shape required for the input tensor.

For models trained for image classification / detection, the shape is always [1, height, width, channels]. To be used as input for RunInference(), this tensor shape must be flattened into a 1-D array with size height * width * channels. To instead get that 1-D array size, use required_input_array_size().

Returns:A 1-D array (numpy.ndarray) representing the required input tensor shape.
get_num_of_output_tensors()

Gets the number of output tensors.

Returns:An integer representing number of output tensors.
get_output_tensor_size(tensor_index)

Gets the size of a specific output tensor.

Parameters:tensor_index (int) – The index position of the output tensor.
Returns:An integer representing the size of the output tensor.
get_raw_output()

Gets the output of the most recent inference.

This can be used by higher level engines for debugging.

Returns:A 1-D array (numpy.ndarray) representing the output tensor. If there are multiple output tensors, they are compressed into a single 1-D array. (Same as what’s returned by RunInference().)
model_path()

Gets the file path for model loaded by this inference engine.

Returns:A string representing the model file’s path.
required_input_array_size()

Gets the required size for the input_tensor given to RunInference().

This is the total size of the 1-D array, once the tensor shape is flattened.

Returns:An integer representing the required input tensor size.
total_output_array_size()
Gets the expected size of the 1-D output array returned by RunInference()
and get_raw_output().
Returns:An integer representing the output tensor size.