manage-ml.py

#!/usr/bin/env python3
"""Firebase Admin SDK ML quickstart example."""

import argparse
from datetime import datetime
from datetime import timezone

from beautifultable import BeautifulTable
import firebase_admin
from firebase_admin import ml


# TODO(user): Configure for your project. (See README.md.)
SERVICE_ACCOUNT_KEY = '/path/to/your/service_account_key.json'
STORAGE_BUCKET = 'your-storage-bucket'

credentials = firebase_admin.credentials.Certificate(SERVICE_ACCOUNT_KEY)
firebase_admin.initialize_app(credentials, options={
    'storageBucket': STORAGE_BUCKET
})


def upload_model(model_file, name, tags=None):
  """Upload a tflite model file to the project and publish it."""
  # Load a tflite file and upload it to Cloud Storage
  print('Uploading to Cloud Storage...')
  model_source = ml.TFLiteGCSModelSource.from_tflite_model_file(model_file)

  # Create the model object
  tflite_format = ml.TFLiteFormat(model_source=model_source)
  model = ml.Model(
      display_name=name,
      model_format=tflite_format)
  if tags is not None:
    model.tags = tags

  # Add the model to your Firebase project and publish it
  new_model = ml.create_model(model)
  ml.publish_model(new_model.model_id)

  print('Model uploaded and published:')
  print_models([new_model], headers=False)


def add_automl_model(model_ref, name, tags=None):
  """Add an AutoML tflite model file to the project and publish it."""
  # Create the model object
  model_source = ml.TFLiteAutoMlSource(model_ref)
  model = ml.Model(
      display_name=name,
      model_format=ml.TFLiteFormat(model_source=model_source))
  if tags is not None:
    model.tags = tags

  # Add the model to your Firebase project and publish it
  new_model = ml.create_model(model)
  new_model.wait_for_unlocked()
  ml.publish_model(new_model.model_id)

  print('Model uploaded and published:')
  print_models([new_model], headers=False)


def list_models(filter_exp=''):
  """List the models in the project."""
  models = ml.list_models(list_filter=filter_exp).iterate_all()
  print_models(models)


def print_models(models, headers=True):
  """Prettyprint a list of models."""
  table = BeautifulTable()
  if headers:
    table.columns.header = ['Name', 'ID', 'Tags']
  for model in models:
    tags = ', '.join(model.tags) if model.tags is not None else ''
    table.rows.append([model.display_name, model.model_id, tags])
  table.set_style(BeautifulTable.STYLE_COMPACT)
  table.columns.header.alignment = BeautifulTable.ALIGN_CENTER
  table.columns.alignment = BeautifulTable.ALIGN_LEFT
  print(table)


def get_model_info(model_id):
  """Get model details."""
  model = ml.get_model(model_id)
  created = datetime.fromtimestamp(model.create_time / 1000, timezone.utc)
  updated = datetime.fromtimestamp(model.update_time / 1000, timezone.utc)
  table = BeautifulTable()
  table.columns.append(['Name:', 'ID:', 'Tags:', 'Published:', 'ETag:',
                        'SHA256:', 'Created:', 'Updated:'])
  table.columns.append([model.display_name,
                        model.model_id,
                        ', '.join(model.tags) if model.tags else '',
                        'Yes' if model.published else 'No',
                        model.etag,
                        model.model_hash,
                        created.isoformat(' ', timespec='seconds'),
                        updated.isoformat(' ', timespec='seconds')])
  table.set_style(BeautifulTable.STYLE_COMPACT)
  table.columns.alignment = BeautifulTable.ALIGN_LEFT
  print(table)


def update_model(model_id, model_file=None, name=None,
                 new_tags=None, remove_tags=None):
  """Update one of the project's models."""
  model = ml.get_model(model_id)

  if model_file is not None:
    # Load a tflite file and upload it to Cloud Storage
    print('Uploading to Cloud Storage...')
    model_source = ml.TFLiteGCSModelSource.from_tflite_model_file(model_file)
    tflite_format = ml.TFLiteFormat(model_source=model_source)
    model.model_format = tflite_format

  if name is not None:
    model.display_name = name

  if new_tags is not None:
    model.tags = new_tags if model.tags is None else model.tags + new_tags

  if remove_tags is not None and model.tags is not None:
    model.tags = list(set(model.tags).difference(set(remove_tags)))

  updated_model = ml.update_model(model)
  ml.publish_model(updated_model.model_id)


def delete_model(model_id):
  """Delete a model from the project."""
  ml.delete_model(model_id)


# The rest of the file just parses the command line and dispatches one of the
# functions above.


def main():
  main_parser = argparse.ArgumentParser()
  subparsers = main_parser.add_subparsers(
      dest='command', required=True, metavar='command')

  new_parser = subparsers.add_parser(
      'new', help='upload a tflite model file or AutoML model reference to'
                  ' your project')
  new_parser.add_argument(
      'name', type=str, help='display name for the new model')
  new_source_group = new_parser.add_mutually_exclusive_group(required=True)
  new_source_group.add_argument(
      '-f', '--file', type=str, help='path to the tflite file')
  new_source_group.add_argument(
      '-a', '--automl', type=str, help='AutoML model reference (e.g. projects/'
                                       '12345678/locations/us-central1/models/'
                                       'ICN1234567890)')
  new_parser.add_argument(
      '-t', '--tags', type=str, help='comma-separated list of tags')

  list_parser = subparsers.add_parser(
      'list', help='list your project\'s models')
  list_parser.add_argument(
      '-f', '--filter', type=str, default='',
      help='''filter expression to limit results (see:
              https://firebase.google.com/docs/ml-kit/manage-hosted-models#list_your_projects_models)''')

  info_parser = subparsers.add_parser(
      'info', help='')
  info_parser.add_argument(
      'model_id', type=valid_id, help='the ID of the model you want to view')

  update_parser = subparsers.add_parser(
      'update', help='update one of your project\'s models')
  update_parser.add_argument(
      'model_id', type=valid_id, help='the ID of the model you want to update')
  update_parser.add_argument(
      '-m', '--model_file', type=str, help='path to a new tflite file')
  update_parser.add_argument(
      '-n', '--name', type=str, help='display name for the model')
  update_parser.add_argument(
      '-t', '--new_tags', type=str,
      help='comma-separated list of tags to add')
  update_parser.add_argument(
      '-d', '--remove_tags', type=str,
      help='comma-separated list of tags to remove')

  delete_parser = subparsers.add_parser(
      'delete', help='delete a model from your project')
  delete_parser.add_argument(
      'model_id', type=valid_id, help='the ID of the model you want to delete')

  args = main_parser.parse_args()
  try:
    if args.command == 'new' and args.file is not None:
      tags = args.tags.split(',') if args.tags is not None else None
      upload_model(args.file.strip(), args.name.strip(), tags)
    if args.command == 'new' and args.automl is not None:
      tags = args.tags.split(',') if args.tags is not None else None
      add_automl_model(args.automl.strip(), args.name.strip(), tags)
    elif args.command == 'list':
      list_models(args.filter)
    elif args.command == 'info':
      get_model_info(args.model_id)
    elif args.command == 'update':
      new_tags = args.new_tags.split(',') if args.new_tags is not None else None
      remove_tags = (
          args.remove_tags.split(',') if args.remove_tags is not None else None)
      update_model(args.model_id, args.model_file, args.name,
                   new_tags, remove_tags)
    elif args.command == 'delete':
      delete_model(args.model_id)
  except firebase_admin.exceptions.NotFoundError:
    print('ERROR: Model not found. Make sure you\'re specifying a valid'
          ' numerical model ID.')


def valid_id(model_id):
  try:
    val = int(model_id)
    return str(val)
  except ValueError:
    raise argparse.ArgumentTypeError('must be a numerical model ID.')


if __name__ == '__main__':
  main()