bitsy-ai · leigh-johnson · Oct 9, 2020 · Oct 9, 2020
diff --git a/MANIFEST.in b/MANIFEST.in
@@ -4,7 +4,7 @@ include HISTORY.md
 include LICENSE
 include README.md
 
-include rpi_deep_pantilt/data/mscoco_label_map.pbtxt
+include rpi_deep_pantilt/data/*
 
 recursive-exclude * __pycache__
 recursive-exclude * *.py[co]

diff --git a/Makefile b/Makefile
@@ -51,7 +51,7 @@ clean-test: ## remove test and coverage artifacts
 	rm -fr .pytest_cache
 
 lint: ## check style with flake8
-	flake8 rpi_deep_pantilt tests
+	black rpi_deep_pantilt tests
 
 test: ## run tests quickly with the default Python
 	pytest
@@ -102,6 +102,12 @@ rpi-deps:
 protoc:
 	cd $$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())') && protoc object_detection/protos/*.proto --python_out=.
 
+edgetpu-deps:
+	echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list
+	curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -
+	sudo apt-get update
+	sudo apt-get install libedgetpu1-std
+
 edgetpu-image:
 	docker build -t  edge_tpu_converter -f tools/edgetpu.Dockerfile .
 
diff --git a/rpi_deep_pantilt/__init__.py b/rpi_deep_pantilt/__init__.py
@@ -3,5 +3,5 @@
 """Top-level package for Raspberry Pi Deep PanTilt."""
 
 __author__ = """Leigh Johnson"""
-__email__ = '[email protected]'
-__version__ = '1.2.1'
+__email__ = "[email protected]"
+__version__ = "__version__ = '2.0.0rc'"
diff --git a/rpi_deep_pantilt/cli.py b/rpi_deep_pantilt/cli.py
@@ -1,149 +1,222 @@
 # -*- coding: utf-8 -*-
 
 """Console script for rpi_deep_pantilt."""
+import importlib
 import logging
+import pprint
 import sys
+
 import click
 
+
 from rpi_deep_pantilt.detect.camera import run_stationary_detect
+from rpi_deep_pantilt.detect.registry import ModelRegistry
+
+# from rpi_deep_pantilt.detect.v1.ssd_mobilenet_v3_coco import (
+#     SSDMobileNet_V3_Small_Coco_PostProcessed,
+#     SSDMobileNet_V3_Coco_EdgeTPU_Quant,
+#     LABELS as SSDMobileNetLabels
+# )
+# from rpi_deep_pantilt.detect.v1.facessd_mobilenet_v2 import (
+#     FaceSSD_MobileNet_V2,
+#     FaceSSD_MobileNet_V2_EdgeTPU,
+#     LABELS as FaceSSDLabels
+# )
 
-from rpi_deep_pantilt.detect.ssd_mobilenet_v3_coco import (
-    SSDMobileNet_V3_Small_Coco_PostProcessed,
-    SSDMobileNet_V3_Coco_EdgeTPU_Quant,
-    LABELS as SSDMobileNetLabels
-)
-from rpi_deep_pantilt.detect.facessd_mobilenet_v2 import (
-    FaceSSD_MobileNet_V2,
-    FaceSSD_MobileNet_V2_EdgeTPU,
-    LABELS as FaceSSDLabels
-)
 from rpi_deep_pantilt.control.manager import pantilt_process_manager
 from rpi_deep_pantilt.control.hardware_test import pantilt_test, camera_test
 
 
-def validate_labels(labels):
-    for label in labels:
-        if label not in (SSDMobileNetLabels + FaceSSDLabels):
-            logging.error(f'''
-            Invalid label: {label} \n
-            Please choose any of the following labels: \n
-            {SSDMobileNetLabels + FaceSSDLabels}
-            ''')
-            sys.exit(1)
-
-
 @click.group()
 def cli():
     pass
 
 
 @cli.command()
-@click.argument('labels', nargs=-1)
-@click.option('--loglevel', required=False, type=str, default='WARNING', help='Run object detection without pan-tilt controls. Pass --loglevel=DEBUG to inspect FPS.')
-@click.option('--edge-tpu', is_flag=True, required=False, type=bool, default=False, help='Accelerate inferences using Coral USB Edge TPU')
-@click.option('--rotation', default=0, type=int, help='PiCamera rotation. If you followed this guide, a rotation value of 0 is correct. https://medium.com/@grepLeigh/real-time-object-tracking-with-tensorflow-raspberry-pi-and-pan-tilt-hat-2aeaef47e134')
-def detect(labels, loglevel, edge_tpu, rotation):
-    '''
-        rpi-deep-pantilt detect [OPTIONS] [LABELS]...
+@click.argument("labels", nargs=-1)
+@click.option(
+    "--api-version",
+    required=False,
+    type=click.Choice(["v1", "v2"]),
+    default="v2",
+    help="API Version to use (default: 2). API v1 is supported for legacy use cases.",
+)
+@click.option(
+    "--predictor",
+    required=False,
+    type=str,
+    default=None,
+    help="Path and module name of a custom predictor class inheriting from rpi_deep_pantilt.detect.custom.base_predictors.BasePredictor",
+)
+@click.option(
+    "--loglevel",
+    required=False,
+    type=str,
+    default="WARNING",
+    help="Run object detection without pan-tilt controls. Pass --loglevel=DEBUG to inspect FPS.",
+)
+@click.option(
+    "--edge-tpu",
+    is_flag=True,
+    required=False,
+    type=bool,
+    default=False,
+    help="Accelerate inferences using Coral USB Edge TPU",
+)
+@click.option(
+    "--rotation",
+    default=0,
+    type=int,
+    help="PiCamera rotation. If you followed this guide, a rotation value of 0 is correct. https://medium.com/@grepLeigh/real-time-object-tracking-with-tensorflow-raspberry-pi-and-pan-tilt-hat-2aeaef47e134",
+)
+@click.option(
+    "--dtype",
+    type=click.Choice(["uint8", "float32"], case_sensitive=True),
+    default="uint8",
+)
+def detect(api_version, labels, predictor, loglevel, edge_tpu, rotation, dtype):
+    """
+    rpi-deep-pantilt detect [OPTIONS] [LABELS]...
 
-        LABELS (optional)
-            One or more labels to detect, for example:
-            $ rpi-deep-pantilt detect person book "wine glass"
+    LABELS (optional)
+        One or more labels to detect, for example:
+        $ rpi-deep-pantilt detect person book "wine glass"
 
-            If no labels are specified, model will detect all labels in this list:
-            $ rpi-deep-pantilt list-labels
+        If no labels are specified, model will detect all labels in this list:
+        $ rpi-deep-pantilt list-labels
 
-        Detect command will automatically load the appropriate model
+    Detect command will automatically load the appropriate model
 
-        For example, providing "face" as the only label will initalize FaceSSD_MobileNet_V2 model
-        $ rpi-deep-pantilt detect face
+    For example, providing "face" as the only label will initalize FaceSSD_MobileNet_V2 model
+    $ rpi-deep-pantilt detect face
 
-        Other labels use SSDMobileNetV3 model with COCO labels
-        $ rpi-deep-pantilt detect person "wine class" orange
-    '''
+    Other labels use SSDMobileNetV3 model with COCO labels
+    $ rpi-deep-pantilt detect person "wine class" orange
+    """
     level = logging.getLevelName(loglevel)
     logging.getLogger().setLevel(level)
 
-    # TypeError: nargs=-1 in combination with a default value is not supported.
-    if not labels:
-        labels = SSDMobileNetLabels
-    # Sanity-check provided labels are supported by model
-    else:
-        validate_labels(labels)
-
-    if 'face' in labels and len(labels) > 1:
-        logging.error(
-            f'''Face detector does not support detection for non-face labels \n
-               Please re-run with face as the only label argument: \n
-               $ rpi-deep-pantilt detect face
-            '''
+    if api_version == "v1" and dtype == "uint8":
+        logging.warning(
+            "WARNING! --dtype=uint8 is not supported by --api-version=v1. Falling back to --dtype=float32."
         )
+        dtype = "float32"
 
-    # FaceSSD model
-    if 'face' in labels:
-        if edge_tpu:
-            model_cls = FaceSSD_MobileNet_V2_EdgeTPU
-        else:
-            model_cls = FaceSSD_MobileNet_V2
-    # All other labels are detected by SSDMobileNetV3 model
+    if predictor is not None:
+        predictor_cls = importlib.import_module(predictor)
     else:
-        if edge_tpu:
-            model_cls = SSDMobileNet_V3_Coco_EdgeTPU_Quant
-        else:
-            model_cls = SSDMobileNet_V3_Small_Coco_PostProcessed
+        # TypeError: nargs=-1 in combination with a default value is not supported.
+        model_registry = ModelRegistry(
+            edge_tpu=edge_tpu, api_version=api_version, dtype=dtype
+        )
+        predictor_cls = model_registry.select_model(labels)
 
-    logging.warning(f'Detecting labels: {labels}')
-    run_stationary_detect(labels, model_cls, rotation)
+    if len(labels) is 0:
+        labels = predictor_cls.LABELS
+
+    run_stationary_detect(labels, predictor_cls, rotation)
 
 
 @cli.command()
-@click.option('--loglevel', required=False, type=str, default='WARNING', help='List all valid classification labels')
-def list_labels(loglevel):
-    level = logging.getLevelName(loglevel)
-    logging.getLogger().setLevel(level)
-    model = SSDMobileNet_V3_Small_Coco_PostProcessed()
-    print('You can detect / track the following objects:')
-    print([x['name'] for x in model.category_index.values()])
+@click.option(
+    "--api-version",
+    required=False,
+    type=int,
+    default=2,
+    help="API Version to use (default: 2). API v1 is supported for legacy use cases.",
+)
+@click.option("--edge-tpu", is_flag=True, required=False, type=bool, default=False)
+@click.option(
+    "--dtype",
+    type=click.Choice(["uint8", "float32"], case_sensitive=True),
+    default="uint8",
+)
+def list_labels(api_version, edge_tpu, dtype):
+    model_registry = ModelRegistry(
+        edge_tpu=edge_tpu, api_version=api_version, dtype=dtype
+    )
+    pp = pprint.PrettyPrinter(indent=1)
+    pp.pprint("The following labels are supported by pretrained models:")
+    pp.pprint(model_registry.label_map())
 
 
 @cli.command()
-@click.argument('label', type=str, default='person')
-@click.option('--loglevel', required=False, type=str, default='WARNING', help='Pass --loglevel=DEBUG to inspect FPS and tracking centroid X/Y coordinates')
-@click.option('--edge-tpu', is_flag=True, required=False, type=bool, default=False, help='Accelerate inferences using Coral USB Edge TPU')
-@click.option('--rotation', default=0, type=int, help='PiCamera rotation. If you followed this guide, a rotation value of 0 is correct. https://medium.com/@grepLeigh/real-time-object-tracking-with-tensorflow-raspberry-pi-and-pan-tilt-hat-2aeaef47e134')
-def track(label, loglevel, edge_tpu, rotation):
-    '''
-        rpi-deep-pantilt track [OPTIONS] [LABEL]
-
-        LABEL (required)
-            Exactly one label to detect, for example:
-            $ rpi-deep-pantilt track person
-
-        Track command will automatically load the appropriate model
-
-        For example, providing "face" will initalize FaceSSD_MobileNet_V2 model
-        $ rpi-deep-pantilt track face
-
-        Other labels use SSDMobileNetV3 model with COCO labels
-        $ rpi-deep-pantilt detect orange
-    '''
+@click.argument("label", type=str, default="person")
+@click.option(
+    "--api-version",
+    required=False,
+    type=click.Choice(["v1", "v2"]),
+    default="v2",
+    help="API Version to use (default: 2). API v1 is supported for legacy use cases.",
+)
+@click.option(
+    "--predictor",
+    required=False,
+    type=str,
+    default=None,
+    help="Path and module name of a custom predictor class inheriting from rpi_deep_pantilt.detect.custom.base_predictors.BasePredictor",
+)
+@click.option(
+    "--loglevel",
+    required=False,
+    type=str,
+    default="WARNING",
+    help="Pass --loglevel=DEBUG to inspect FPS and tracking centroid X/Y coordinates",
+)
+@click.option(
+    "--edge-tpu",
+    is_flag=True,
+    required=False,
+    type=bool,
+    default=False,
+    help="Accelerate inferences using Coral USB Edge TPU",
+)
+@click.option(
+    "--rotation",
+    default=0,
+    type=int,
+    help="PiCamera rotation. If you followed this guide, a rotation value of 0 is correct. https://medium.com/@grepLeigh/real-time-object-tracking-with-tensorflow-raspberry-pi-and-pan-tilt-hat-2aeaef47e134",
+)
+@click.option(
+    "--dtype",
+    type=click.Choice(["uint8", "float32"], case_sensitive=True),
+    default="uint8",
+)
+def track(label, api_version, predictor, loglevel, edge_tpu, rotation, dtype):
+    """
+    rpi-deep-pantilt track [OPTIONS] [LABEL]
+
+    LABEL (required)
+        Exactly one label to detect, for example:
+        $ rpi-deep-pantilt track person
+
+    Track command will automatically load the appropriate model
+
+    For example, providing "face" will initalize FaceSSD_MobileNet_V2 model
+    $ rpi-deep-pantilt track face
+
+    Other labels use SSDMobileNetV3 model with COCO labels
+    $ rpi-deep-pantilt detect orange
+    """
     level = logging.getLevelName(loglevel)
     logging.getLogger().setLevel(level)
 
-    validate_labels((label,))
+    if api_version == "v1" and dtype == "uint8":
+        logging.warning(
+            "WARNING! --dtype=uint8 is not supported by --api-version=v1. Falling back to --dtype=float32."
+        )
+        dtype = "float32"
 
-    if label == 'face':
-        if edge_tpu:
-            model_cls = FaceSSD_MobileNet_V2_EdgeTPU
-        else:
-            model_cls = FaceSSD_MobileNet_V2
+    if predictor is not None:
+        predictor_cls = importlib.import_module(predictor)
     else:
-        if edge_tpu:
-            model_cls = SSDMobileNet_V3_Coco_EdgeTPU_Quant
-        else:
-            model_cls = SSDMobileNet_V3_Small_Coco_PostProcessed
+        # TypeError: nargs=-1 in combination with a default value is not supported.
+        model_registry = ModelRegistry(
+            edge_tpu=edge_tpu, api_version=api_version, dtype=dtype
+        )
+        predictor_cls = model_registry.select_model((label,))
 
-    return pantilt_process_manager(model_cls, labels=(label,),  rotation=rotation)
+    return pantilt_process_manager(predictor_cls, labels=(label,), rotation=rotation)
 
 
 @cli.group()
@@ -152,16 +225,21 @@ def test():
 
 
 @test.command()
-@click.option('--loglevel', required=False, type=str, default='INFO')
+@click.option("--loglevel", required=False, type=str, default="INFO")
 def pantilt(loglevel):
     level = logging.getLevelName(loglevel)
     logging.getLogger().setLevel(level)
     return pantilt_test()
 
 
 @test.command()
-@click.option('--loglevel', required=False, type=str, default='INFO')
-@click.option('--rotation', default=0, type=int, help='PiCamera rotation. If you followed this guide, a rotation value of 0 is correct. https://medium.com/@grepLeigh/real-time-object-tracking-with-tensorflow-raspberry-pi-and-pan-tilt-hat-2aeaef47e134')
+@click.option("--loglevel", required=False, type=str, default="INFO")
+@click.option(
+    "--rotation",
+    default=0,
+    type=int,
+    help="PiCamera rotation. If you followed this guide, a rotation value of 0 is correct. https://medium.com/@grepLeigh/real-time-object-tracking-with-tensorflow-raspberry-pi-and-pan-tilt-hat-2aeaef47e134",
+)
 def camera(loglevel, rotation):
     level = logging.getLevelName(loglevel)
     logging.getLogger().setLevel(level)