#!/usr/bin/env python3
# Copyright 2017 Google Inc.
#
# 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.
"""Script to run generic MobileNet based classification model."""
import argparse

from picamera import PiCamera

from aiy.vision import inference
from aiy.vision.models import utils


def read_labels(label_path):
  with open(label_path) as label_file:
    return [label.strip() for label in label_file.readlines()]


def process(result, labels, out_tensor_name, threshold, top_k):
  """Processes inference result and returns labels sorted by confidence."""
  # MobileNet based classification model returns one result vector.
  assert len(result.tensors) == 1
  tensor = result.tensors[out_tensor_name]
  probs, shape = tensor.data, tensor.shape
  assert shape.depth == len(labels)
  pairs = [pair for pair in enumerate(probs) if pair[1] > threshold]
  pairs = sorted(pairs, key=lambda pair: pair[1], reverse=True)
  pairs = pairs[0:top_k]
  return [' %s (%.2f)' % (labels[index], prob) for index, prob in pairs]


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--model_path',
      required=True,
      help='Path to converted model file that can run on VisionKit.')
  parser.add_argument(
      '--label_path',
      required=True,
      help='Path to label file that corresponds to the model.')
  parser.add_argument(
      '--input_height', type=int, required=True, help='Input height.')
  parser.add_argument(
      '--input_width', type=int, required=True, help='Input width.')
  parser.add_argument(
      '--input_layer', required=True, help='Name of input layer.')
  parser.add_argument(
      '--output_layer', required=True, help='Name of output layer.')
  parser.add_argument(
      '--num_frames',
      type=int,
      default=-1,
      help='Sets the number of frames to run for, otherwise runs forever.')
  parser.add_argument(
      '--input_mean', type=float, default=128.0, help='Input mean.')
  parser.add_argument(
      '--input_std', type=float, default=128.0, help='Input std.')
  parser.add_argument(
      '--input_depth', type=int, default=3, help='Input depth.')
  parser.add_argument(
      '--threshold', type=float, default=0.1,
      help='Threshold for classification score (from output tensor).')
  parser.add_argument(
      '--top_k', type=int, default=3, help='Keep at most top_k labels.')
  args = parser.parse_args()

  model = inference.ModelDescriptor(
      name='mobilenet_based_classifier',
      input_shape=(1, args.input_height, args.input_width, args.input_depth),
      input_normalizer=(args.input_mean, args.input_std),
      compute_graph=utils.load_compute_graph(args.model_path))
  labels = read_labels(args.label_path)

  with PiCamera(sensor_mode=4, resolution=(1640, 1232), framerate=30) as camera:
    with inference.CameraInference(model) as camera_inference:
      for i, result in enumerate(camera_inference.run()):
        if i == args.num_frames:
          break
        processed_result = process(result, labels, args.output_layer,
                                   args.threshold, args.top_k)
        if processed_result:
          print('Detecting:')
          print('\n'.join(processed_result))
        else:
          print('Nothing detected with threshold=%.2f and top_k=%d' % (args.threshold, args.top_k))


if __name__ == '__main__':
  main()

Add script to run generic MobileNet based classification model.