roughly measuring the "switch" time between 2 models

switch.py

import threading
import time
import cv2
import cv2 as cv

import tensorflow as tf
from tensorflow.lite.python.interpreter import load_delegate
from edgetpu.basic import edgetpu_utils
from collections import namedtuple
import numpy as np

def get_input_tensor(interpreter, image_path):
    input_details = interpreter.get_input_details()
    _, height, width, _ = input_details[0]['shape']
    input = cv.imread(image_path, cv2.IMREAD_COLOR)
    input = cv.resize(input, (width, height)).astype(np.uint8)
    return np.reshape(input, newshape=[1, width, height, 3], order='C')

def run_one_model_one_tpu(model_path, image_path, num_inferences):
    start_time = time.perf_counter()
    interpreter = tf.lite.Interpreter(model_path=model_path,
                                            experimental_delegates=[load_delegate('libedgetpu.so.1.0')])
    interpreter.allocate_tensors()
    input_tensor = get_input_tensor(interpreter, image_path=image_path)
    interpreter.set_tensor(interpreter.get_input_details()[0]["index"], input_tensor)

    for _ in range(num_inferences):
        interpreter.invoke()
    return time.perf_counter() - start_time

def run_two_models_one_tpu(first_model_path, second_model_path, image_path1, image_path2, num_inferences):
    start_time = time.perf_counter()
    interpreter_first = tf.lite.Interpreter(model_path=first_model_path,
                                    experimental_delegates=[load_delegate('libedgetpu.so.1.0')])
    interpreter_first.allocate_tensors()
    interpreter_second = tf.lite.Interpreter(model_path=second_model_path,
                                          experimental_delegates=[load_delegate('libedgetpu.so.1.0')])
    interpreter_second.allocate_tensors()
    input_tensor_first_model = get_input_tensor(interpreter_first, image_path=image_path1)
    input_tensor_second_model = get_input_tensor(interpreter_second, image_path=image_path2)
    interpreter_first.set_tensor(interpreter_first.get_input_details()[0]["index"], input_tensor_first_model)
    interpreter_second.set_tensor(interpreter_second.get_input_details()[0]["index"], input_tensor_second_model)

    for _ in range(num_inferences):
        interpreter_first.invoke()
        interpreter_second.invoke()
    return time.perf_counter() - start_time

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--first_model',
        help='Path of first model.',
        default='/tmp/test_data/inat_plant_edgetpu.tflite')
    parser.add_argument(
        '--second_model',
        help='Path of second model.',
        default='/tmp/test_data/inat_bird_edgetpu.tflite')
    parser.add_argument(
        '--image_path1',
        help='image 1.',
        default='/tmp/test_data/plant.bmp')
    parser.add_argument(
        '--image_path2',
        help='image2.',
        default='/tmp/test_data/bird.bmp')

    args = parser.parse_args()

    # TPU runtime version
    tpu_runtime_ver = edgetpu_utils.GetRuntimeVersion()
    print('TPU runtime version:'+ tpu_runtime_ver)
    edge_tpus = edgetpu_utils.ListEdgeTpuPaths(
        edgetpu_utils.EDGE_TPU_STATE_UNASSIGNED)

    print('Number of available TPU = %d'%(len(edge_tpus)))

    for num_inferences in [100, 200, 400, 800, 1600, 3200, 6400]:
        print('\n-----num inference: %d-----' % num_inferences)
        costs_first_model = run_one_model_one_tpu(args.first_model, args.image_path1, num_inferences)
        print('Cost for first model: %.2f seconds.' % costs_first_model)
        costs_second_model = run_one_model_one_tpu(args.second_model, args.image_path2, num_inferences)
        print('Cost for second model: %.2f seconds.' % costs_second_model)
        cost_two_models = run_two_models_one_tpu(args.first_model, args.second_model,
            args.image_path1, args.image_path2, num_inferences)
        print('Cost for both models: %.2f seconds' % cost_two_models)
        print('Cost of switching model %.2f seconds.' % (cost_two_models - (costs_first_model + costs_second_model)))


if __name__ == '__main__':
    main()