diff --git a/MANIFEST.in b/MANIFEST.in index 8ddf2c6..5ce97c8 100644 --- 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 index 84f71a3..9bbe23c 100644 --- 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 index ec88ef0..c515e24 100644 --- 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__ = 'hi@leighjohnson.me' -__version__ = '1.2.1' +__email__ = "hi@leighjohnson.me" +__version__ = "__version__ = '2.0.0rc'" diff --git a/rpi_deep_pantilt/cli.py b/rpi_deep_pantilt/cli.py index 0cb5361..eb6823e 100644 --- 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,7 +225,7 @@ 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) @@ -160,8 +233,13 @@ def pantilt(loglevel): @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) diff --git a/rpi_deep_pantilt/control/hardware_test.py b/rpi_deep_pantilt/control/hardware_test.py index e84693c..0972b1c 100644 --- a/rpi_deep_pantilt/control/hardware_test.py +++ b/rpi_deep_pantilt/control/hardware_test.py @@ -12,20 +12,20 @@ def camera_test(rotation): camera = PiCamera() camera.rotation = rotation - logging.info('Starting Raspberry Pi Camera') + logging.info("Starting Raspberry Pi Camera") camera.start_preview() try: while True: continue except KeyboardInterrupt: - logging.info('Stopping Raspberry Pi Camera') + logging.info("Stopping Raspberry Pi Camera") camera.stop_preview() def pantilt_test(): - logging.info('Starting Pan-Tilt HAT test!') - logging.info('Pan-Tilt HAT should follow a smooth sine wave') + logging.info("Starting Pan-Tilt HAT test!") + logging.info("Pan-Tilt HAT should follow a smooth sine wave") while True: # Get the time in seconds t = time.time() @@ -41,5 +41,5 @@ def pantilt_test(): time.sleep(0.005) -if __name__ == '__main__': +if __name__ == "__main__": pantilt_test() diff --git a/rpi_deep_pantilt/control/manager.py b/rpi_deep_pantilt/control/manager.py index 190a81b..5700bda 100644 --- a/rpi_deep_pantilt/control/manager.py +++ b/rpi_deep_pantilt/control/manager.py @@ -16,10 +16,7 @@ SERVO_MIN = -90 SERVO_MAX = 90 -CENTER = ( - RESOLUTION[0] // 2, - RESOLUTION[1] // 2 -) +CENTER = (RESOLUTION[0] // 2, RESOLUTION[1] // 2) # function to handle keyboard interrupt @@ -38,7 +35,7 @@ def signal_handler(sig, frame): def in_range(val, start, end): # determine the input vale is in the supplied range - return (val >= start and val <= end) + return val >= start and val <= end def set_servos(pan, tilt): @@ -53,12 +50,12 @@ def set_servos(pan, tilt): if in_range(pan_angle, SERVO_MIN, SERVO_MAX): pth.pan(pan_angle) else: - logging.info(f'pan_angle not in range {pan_angle}') + logging.info(f"pan_angle not in range {pan_angle}") if in_range(tilt_angle, SERVO_MIN, SERVO_MAX): pth.tilt(tilt_angle) else: - logging.info(f'tilt_angle not in range {tilt_angle}') + logging.info(f"tilt_angle not in range {tilt_angle}") def pid_process(output, p, i, d, box_coord, origin_coord, action): @@ -75,51 +72,54 @@ def pid_process(output, p, i, d, box_coord, origin_coord, action): output.value = p.update(error) # logging.info(f'{action} error {error} angle: {output.value}') + # ('person',) -#('orange', 'apple', 'sports ball') +# ('orange', 'apple', 'sports ball') -def pantilt_process_manager( - model_cls, - labels=('person',), - rotation=0 -): +def pantilt_process_manager(model_cls, labels=("person",), rotation=0): pth.servo_enable(1, True) pth.servo_enable(2, True) with Manager() as manager: # set initial bounding box (x, y)-coordinates to center of frame - center_x = manager.Value('i', 0) - center_y = manager.Value('i', 0) + center_x = manager.Value("i", 0) + center_y = manager.Value("i", 0) center_x.value = RESOLUTION[0] // 2 center_y.value = RESOLUTION[1] // 2 # pan and tilt angles updated by independent PID processes - pan = manager.Value('i', 0) - tilt = manager.Value('i', 0) + pan = manager.Value("i", 0) + tilt = manager.Value("i", 0) # PID gains for panning - pan_p = manager.Value('f', 0.05) + pan_p = manager.Value("f", 0.05) # 0 time integral gain until inferencing is faster than ~50ms - pan_i = manager.Value('f', 0.1) - pan_d = manager.Value('f', 0) + pan_i = manager.Value("f", 0.1) + pan_d = manager.Value("f", 0) # PID gains for tilting - tilt_p = manager.Value('f', 0.15) + tilt_p = manager.Value("f", 0.15) # 0 time integral gain until inferencing is faster than ~50ms - tilt_i = manager.Value('f', 0.2) - tilt_d = manager.Value('f', 0) + tilt_i = manager.Value("f", 0.2) + tilt_d = manager.Value("f", 0) - detect_processr = Process(target=run_pantilt_detect, - args=(center_x, center_y, labels, model_cls, rotation)) + detect_processr = Process( + target=run_pantilt_detect, + args=(center_x, center_y, labels, model_cls, rotation), + ) - pan_process = Process(target=pid_process, - args=(pan, pan_p, pan_i, pan_d, center_x, CENTER[0], 'pan')) + pan_process = Process( + target=pid_process, + args=(pan, pan_p, pan_i, pan_d, center_x, CENTER[0], "pan"), + ) - tilt_process = Process(target=pid_process, - args=(tilt, tilt_p, tilt_i, tilt_d, center_y, CENTER[1], 'tilt')) + tilt_process = Process( + target=pid_process, + args=(tilt, tilt_p, tilt_i, tilt_d, center_y, CENTER[1], "tilt"), + ) servo_process = Process(target=set_servos, args=(pan, tilt)) @@ -134,5 +134,5 @@ def pantilt_process_manager( servo_process.join() -if __name__ == '__main__': +if __name__ == "__main__": pantilt_process_manager() diff --git a/rpi_deep_pantilt/control/pid.py b/rpi_deep_pantilt/control/pid.py index 0dd2b6e..97c2a8b 100644 --- a/rpi_deep_pantilt/control/pid.py +++ b/rpi_deep_pantilt/control/pid.py @@ -45,8 +45,4 @@ def update(self, error, sleep=0.01): self.error_prev = error # sum the terms and return - return sum([ - self.kP * self.cP, - self.kI * self.cI, - self.kD * self.cD] - ) + return sum([self.kP * self.cP, self.kI * self.cI, self.kD * self.cD]) diff --git a/rpi_deep_pantilt/data/leopard.tflite b/rpi_deep_pantilt/data/leopard.tflite new file mode 100644 index 0000000..819175e Binary files /dev/null and b/rpi_deep_pantilt/data/leopard.tflite differ diff --git a/rpi_deep_pantilt/data/leopard_label_map.pbtxt b/rpi_deep_pantilt/data/leopard_label_map.pbtxt new file mode 100644 index 0000000..498acd2 --- /dev/null +++ b/rpi_deep_pantilt/data/leopard_label_map.pbtxt @@ -0,0 +1,6 @@ +item { + name: "leopard" + id: 1 + display_name: "leopard" +} + diff --git a/rpi_deep_pantilt/detect/camera.py b/rpi_deep_pantilt/detect/camera.py index 7d2c514..33740a0 100644 --- a/rpi_deep_pantilt/detect/camera.py +++ b/rpi_deep_pantilt/detect/camera.py @@ -18,11 +18,13 @@ # https://github.com/dtreskunov/rpi-sensorium/commit/40c6f3646931bf0735c5fe4579fa89947e96aed7 -def run_pantilt_detect(center_x, center_y, labels, model_cls, rotation, resolution=RESOLUTION): - ''' - Updates center_x and center_y coordinates with centroid of detected class's bounding box - Overlay is only rendered around the tracked object - ''' +def run_pantilt_detect( + center_x, center_y, labels, model_cls, rotation, resolution=RESOLUTION +): + """ + Updates center_x and center_y coordinates with centroid of detected class's bounding box + Overlay is only rendered around the tracked object + """ model = model_cls() capture_manager = PiCameraStream(resolution=resolution, rotation=rotation) @@ -37,46 +39,48 @@ def run_pantilt_detect(center_x, center_y, labels, model_cls, rotation, resoluti frame = capture_manager.read() prediction = model.predict(frame) - if not len(prediction.get('detection_boxes')): + if not len(prediction.get("detection_boxes")): continue - if any(item in label_idxs for item in prediction.get('detection_classes')): + if any(item in label_idxs for item in prediction.get("detection_classes")): tracked = ( - (i, x) for i, x in - enumerate(prediction.get('detection_classes')) + (i, x) + for i, x in enumerate(prediction.get("detection_classes")) if x in label_idxs ) tracked_idxs, tracked_classes = zip(*tracked) - track_target = prediction.get('detection_boxes')[ - tracked_idxs[0]] + track_target = prediction.get("detection_boxes")[tracked_idxs[0]] # [ymin, xmin, ymax, xmax] y = int( - RESOLUTION[1] - ((np.take(track_target, [0, 2])).mean() * RESOLUTION[1])) + RESOLUTION[1] + - ((np.take(track_target, [0, 2])).mean() * RESOLUTION[1]) + ) center_y.value = y x = int( - RESOLUTION[0] - ((np.take(track_target, [1, 3])).mean() * RESOLUTION[0])) + RESOLUTION[0] + - ((np.take(track_target, [1, 3])).mean() * RESOLUTION[0]) + ) center_x.value = x - display_name = model.category_index[tracked_classes[0]]['name'] - logging.info( - f'Tracking {display_name} center_x {x} center_y {y}') + display_name = model.category_index[tracked_classes[0]]["name"] + logging.info(f"Tracking {display_name} center_x {x} center_y {y}") overlay = model.create_overlay(frame, prediction) capture_manager.overlay_buff = overlay if LOGLEVEL is logging.DEBUG and (time.time() - start_time) > 1: fps_counter += 1 fps = fps_counter / (time.time() - start_time) - logging.debug(f'FPS: {fps}') + logging.debug(f"FPS: {fps}") fps_counter = 0 start_time = time.time() def run_stationary_detect(labels, model_cls, rotation): - ''' - Overlay is rendered around all tracked objects - ''' + """ + Overlay is rendered around all tracked objects + """ model = model_cls() capture_manager = PiCameraStream(resolution=RESOLUTION, rotation=rotation) @@ -93,15 +97,18 @@ def run_stationary_detect(labels, model_cls, rotation): frame = capture_manager.read() prediction = model.predict(frame) - if not len(prediction.get('detection_boxes')): + if not len(prediction.get("detection_boxes")): continue - if any(item in label_idxs for item in prediction.get('detection_classes')): - + if any( + item in label_idxs for item in prediction.get("detection_classes") + ): + # Not all models will need to implement a filter_tracked() interface # For example, FaceSSD only allows you to track 1 class (faces) and does not implement this method try: filtered_prediction = model.filter_tracked( - prediction, label_idxs) + prediction, label_idxs + ) except AttributeError: filtered_prediction = prediction @@ -111,7 +118,7 @@ def run_stationary_detect(labels, model_cls, rotation): if LOGLEVEL is logging.DEBUG and (time.time() - start_time) > 1: fps_counter += 1 fps = fps_counter / (time.time() - start_time) - logging.debug(f'FPS: {fps}') + logging.debug(f"FPS: {fps}") fps_counter = 0 start_time = time.time() except KeyboardInterrupt: @@ -135,24 +142,25 @@ def silent_send_buffer(zelf, *args, **kwargs): class PiCameraStream(object): """ - Continuously capture video frames, and optionally render with an overlay + Continuously capture video frames, and optionally render with an overlay - Arguments - resolution - tuple (x, y) size - framerate - int - vflip - reflect capture on x-axis - hflip - reflect capture on y-axis + Arguments + resolution - tuple (x, y) size + framerate - int + vflip - reflect capture on x-axis + hflip - reflect capture on y-axis """ - def __init__(self, - resolution=(320, 240), - framerate=24, - vflip=False, - hflip=False, - rotation=0, - max_workers=2 - ): + def __init__( + self, + resolution=(320, 240), + framerate=24, + vflip=False, + hflip=False, + rotation=0, + max_workers=2, + ): self.camera = PiCamera() self.camera.resolution = resolution @@ -171,7 +179,7 @@ def __init__(self, self.overlay_buff = None self.frame = None self.stopped = False - logging.info('starting camera preview') + logging.info("starting camera preview") self.camera.start_preview() def render_overlay(self): @@ -180,7 +188,8 @@ def render_overlay(self): self.overlay.update(self.overlay_buff) elif not self.overlay and self.overlay_buff: self.overlay = self.camera.add_overlay( - self.overlay_buff, layer=3, size=self.camera.resolution) + self.overlay_buff, layer=3, size=self.camera.resolution + ) _monkey_patch_picamera(self.overlay) def start_overlay(self): @@ -188,7 +197,7 @@ def start_overlay(self): return self def start(self): - '''Begin handling frame stream in a separate thread''' + """Begin handling frame stream in a separate thread""" Thread(target=self.flush, args=()).start() return self diff --git a/rpi_deep_pantilt/detect/util/LICENSE b/rpi_deep_pantilt/detect/custom/__init__.py similarity index 100% rename from rpi_deep_pantilt/detect/util/LICENSE rename to rpi_deep_pantilt/detect/custom/__init__.py diff --git a/rpi_deep_pantilt/detect/custom/base_predictors.py b/rpi_deep_pantilt/detect/custom/base_predictors.py new file mode 100644 index 0000000..0c56866 --- /dev/null +++ b/rpi_deep_pantilt/detect/custom/base_predictors.py @@ -0,0 +1,221 @@ +# Python +from abc import ABCMeta, abstractmethod +import logging +import os +import sys + + +# lib +import numpy as np +from PIL import Image + +import tensorflow as tf + +from rpi_deep_pantilt.detect.util.label import create_category_index_from_labelmap +from rpi_deep_pantilt.detect.util.visualization import ( + visualize_boxes_and_labels_on_image_array, +) + + +class BasePredictor(metaclass=ABCMeta): + + EDGETPU_SHARED_LIB = "libedgetpu.so.1" + LABELS = [] + + def __init__( + self, + model_name, + tflite_file, + label_file, + model_uri=None, + input_shape=(320, 320), + input_type=tf.uint8, + edge_tpu=False, + min_score_thresh=0.50, + ): + + self.model_uri = model_uri + self.model_name = model_name + self.tflite_file = tflite_file + self.input_type = input_type + self.label_file = label_file + self.min_score_thresh = min_score_thresh + + if model_uri is not None: + self.model_dir = tf.keras.utils.get_file( + fname=self.model_name, + origin=self.model_uri, + untar=True, + cache_subdir="models", + ) + logging.info(f"Downloaded {model_name} to {self.model_dir}") + self.model_path = ( + os.path.splitext(os.path.splitext(self.model_dir)[0])[0] + + f"/{self.tflite_file}" + ) + else: + logging.info(f"Loading {model_name} from {tflite_file}") + self.model_path = tflite_file + + if edge_tpu: + try: + logging.warning("Loading Coral tflite_runtime for Edge TPU") + from tflite_runtime import interpreter as coral_tflite_interpreter + + self.tflite_interpreter = coral_tflite_interpreter.Interpreter( + model_path=self.model_path, + experimental_delegates=[ + tf.lite.experimental.load_delegate(self.EDGETPU_SHARED_LIB) + ], + ) + except ImportError as e: + logging.warning( + "Failed to import Coral Edge TPU tflite_runtime. Falling back to TensorFlow tflite runtime. If you are using an Edge TPU, please run: \n" + ) + logging.warning( + "$ pip install https://dl.google.com/coral/python/tflite_runtime-2.1.0.post1-cp37-cp37m-linux_armv7l.whl" + ) + self.tflite_interpreter = tf.lite.Interpreter( + model_path=self.model_path, + ) + else: + self.tflite_interpreter = tf.lite.Interpreter( + model_path=self.model_path, + ) + + self.tflite_interpreter.allocate_tensors() + self.input_details = self.tflite_interpreter.get_input_details() + self.output_details = self.tflite_interpreter.get_output_details() + + self.category_index = create_category_index_from_labelmap( + label_file, use_display_name=True + ) + + logging.info(f"loaded labels from {self.label_file} \n {self.category_index}") + + logging.info(f"initialized model {model_name} \n") + logging.info(f"model inputs: {self.input_details} \n {self.input_details}") + logging.info(f"model outputs: {self.output_details} \n {self.output_details}") + + def label_to_category_index(self, labels): + return tuple( + map( + lambda x: x["id"], + filter(lambda x: x["name"] in labels, self.category_index.values()), + ) + ) + + def label_display_name_by_idx(self, idx): + return self.category_index[idx]["display_name"] + + @abstractmethod + def create_overlay(self, image_np, output_dict): + pass + + @abstractmethod + def predict(self, image): + pass + + @classmethod + def validate_labels(cls, labels): + return all([x in cls.LABELS for x in labels]) + + +# TFLite_Detection_PostProcess custom op is a non-max supression op (NMS) +# utilized in TensorFlow's Object Detection API / Model zoo +# https://github.com/tensorflow/models/tree/master/research/object_detection +# TFLite_Detection_PostProcess custom op node has four outputs: +# detection_boxes: a float32 tensor of shape [1, num_boxes, 4] with box +# locations +# detection_classes: a float32 tensor of shape [1, num_boxes] +# with class indices +# detection_scores: a float32 tensor of shape [1, num_boxes] +# with class scores +# num_boxes: a float32 tensor of size 1 containing the number of detected +# boxes + +# Without the PostProcessing ops, the graph has two outputs: +# 'raw_outputs/box_encodings': a float32 tensor of shape [1, num_anchors, 4] +# containing the encoded box predictions. +# 'raw_outputs/class_predictions': a float32 tensor of shape +# [1, num_anchors, num_classes] containing the class scores for each anchor +# after applying score conversion. + + +class TFLiteDetectionPostProcessOverlay(BasePredictor): + def __init__(self, *args, max_boxes_to_draw=3, **kwargs): + + self.max_boxes_to_draw = max_boxes_to_draw + super().__init__(*args, **kwargs) + + def create_overlay(self, image_np, output_dict): + + image_np = image_np.copy() + + # draw bounding boxes + visualize_boxes_and_labels_on_image_array( + image_np, + output_dict["detection_boxes"], + output_dict["detection_classes"], + output_dict["detection_scores"], + self.category_index, + use_normalized_coordinates=True, + line_thickness=4, + min_score_thresh=self.min_score_thresh, + max_boxes_to_draw=self.max_boxes_to_draw, + ) + + img = Image.fromarray(image_np) + + return img.tobytes() + + +class TFLiteDetectionPostProcessPredictor(TFLiteDetectionPostProcessOverlay): + def __init__(self, *args, **kwargs): + + super().__init__(*args, **kwargs) + + # float16 is not a supported data type (yet) + assert self.input_type is tf.float32 or self.input_type is tf.uint8 + + def predict(self, image): + + image = np.asarray(image) + # normalize 0 - 255 RGB to values between (-1, 1) if float32 data type + if self.input_type is tf.float32: + image = (image / 128.0) - 1 + # The input needs to be a tensor, convert it using `tf.convert_to_tensor`. + input_tensor = tf.convert_to_tensor(image, dtype=self.input_type) + + # The model expects a batch of images, so add an axis with `tf.newaxis`. + input_tensor = input_tensor[tf.newaxis, ...] + + # Run inference + self.tflite_interpreter.set_tensor(self.input_details[0]["index"], input_tensor) + + self.tflite_interpreter.invoke() + + box_data = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[0]["index"]) + ) + class_data = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[1]["index"]) + ) + score_data = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[2]["index"]) + ) + num_detections = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[3]["index"]) + ) + + # hilarious, but it seems like all classes predictions are off by 1 idx + class_data = tf.squeeze(class_data, axis=[0]).numpy().astype(np.int64) + 1 + box_data = tf.squeeze(box_data, axis=[0]).numpy() + score_data = tf.squeeze(score_data, axis=[0]).numpy() + + return { + "detection_boxes": box_data, + "detection_classes": class_data, + "detection_scores": score_data, + "num_detections": len(num_detections), + } diff --git a/rpi_deep_pantilt/detect/pretrained/__init__.py b/rpi_deep_pantilt/detect/pretrained/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/rpi_deep_pantilt/detect/pretrained/api_v1/__init__.py b/rpi_deep_pantilt/detect/pretrained/api_v1/__init__.py new file mode 100644 index 0000000..06468fd --- /dev/null +++ b/rpi_deep_pantilt/detect/pretrained/api_v1/__init__.py @@ -0,0 +1,45 @@ +from rpi_deep_pantilt.detect.pretrained.api_v1.facessd_mobilenet_v2 import ( + FaceSSDMobileNetV2EdgeTPU, + FaceSSDMobileNetV2Float32, +) + +from rpi_deep_pantilt.detect.pretrained.api_v1.ssd_mobilenet_v3_coco import ( + SSDMobileNetV3EdgeTPU, + SSDMobileNetV3Float32, +) + + +class FaceSSDMobileNetV2Int8(object): + LABELS = [] + + def validate_labels(self): + return False + + def __init__(self, *args, **kwargs): + raise NotImplementedError( + "Please specify --api-version=2 to load FaceSSDMobileNetV2Int8 predictor" + ) + + +class SSDMobileNetV3Int8(object): + LABELS = [] + + def validate_labels(self): + return False + + def __init__(self, *args, **kwargs): + raise NotImplementedError( + "Please specify --api-version=2 to load SSDMobileNetV3Int8 predictor" + ) + + +class LeopardAutoMLInt8(object): + LABELS = [] + + def validate_labels(self): + return False + + def __init__(self, *args, **kwargs): + raise NotImplementedError( + "Please specify--api-version=2 to load LeopardAutoMLInt8 predictor" + ) diff --git a/rpi_deep_pantilt/detect/facessd_mobilenet_v2.py b/rpi_deep_pantilt/detect/pretrained/api_v1/facessd_mobilenet_v2.py similarity index 53% rename from rpi_deep_pantilt/detect/facessd_mobilenet_v2.py rename to rpi_deep_pantilt/detect/pretrained/api_v1/facessd_mobilenet_v2.py index 6aab402..77059c4 100644 --- a/rpi_deep_pantilt/detect/facessd_mobilenet_v2.py +++ b/rpi_deep_pantilt/detect/pretrained/api_v1/facessd_mobilenet_v2.py @@ -10,29 +10,31 @@ import tensorflow as tf from rpi_deep_pantilt import __path__ as rpi_deep_pantilt_path +from rpi_deep_pantilt.detect.util.exceptions import InvalidLabelException from rpi_deep_pantilt.detect.util.label import create_category_index_from_labelmap -from rpi_deep_pantilt.detect.util.visualization import visualize_boxes_and_labels_on_image_array +from rpi_deep_pantilt.detect.util.visualization import ( + visualize_boxes_and_labels_on_image_array, +) -LABELS = ['face'] +class FaceSSDMobileNetV2EdgeTPU(object): -class FaceSSD_MobileNet_V2_EdgeTPU(object): - - EDGETPU_SHARED_LIB = 'libedgetpu.so.1' - PATH_TO_LABELS = rpi_deep_pantilt_path[0] + '/data/facessd_label_map.pbtxt' + EDGETPU_SHARED_LIB = "libedgetpu.so.1" + PATH_TO_LABELS = rpi_deep_pantilt_path[0] + "/data/facessd_label_map.pbtxt" + LABELS = ["face"] def __init__( self, - base_url='https://github.com/leigh-johnson/rpi-deep-pantilt/releases/download/v1.0.1/', - model_name='facessd_mobilenet_v2_quantized_320x320_open_image_v4_tflite2', + base_url="https://github.com/leigh-johnson/rpi-deep-pantilt/releases/download/v1.0.1/", + model_name="facessd_mobilenet_v2_quantized_320x320_open_image_v4_tflite2", input_shape=(320, 320), min_score_thresh=0.50, - tflite_model_file='model_postprocessed_quantized_128_uint8_edgetpu.tflite' + tflite_model_file="model_postprocessed_quantized_128_uint8_edgetpu.tflite", ): self.base_url = base_url self.model_name = model_name - self.model_file = model_name + '.tar.gz' + self.model_file = model_name + ".tar.gz" self.model_url = base_url + self.model_file self.tflite_model_file = tflite_model_file @@ -40,29 +42,31 @@ def __init__( fname=self.model_file, origin=self.model_url, untar=True, - cache_subdir='models' + cache_subdir="models", ) self.min_score_thresh = min_score_thresh - self.model_path = os.path.splitext( - os.path.splitext(self.model_dir)[0] - )[0] + f'/{self.tflite_model_file}' + self.model_path = ( + os.path.splitext(os.path.splitext(self.model_dir)[0])[0] + + f"/{self.tflite_model_file}" + ) try: from tflite_runtime import interpreter as coral_tflite_interpreter except ImportError as e: logging.error(e) - logging.error('Please install Edge TPU tflite_runtime:') + logging.error("Please install Edge TPU tflite_runtime:") logging.error( - '$ pip install https://dl.google.com/coral/python/tflite_runtime-2.1.0.post1-cp37-cp37m-linux_armv7l.whl') + "$ pip install https://dl.google.com/coral/python/tflite_runtime-2.1.0.post1-cp37-cp37m-linux_armv7l.whl" + ) sys.exit(1) self.tflite_interpreter = coral_tflite_interpreter.Interpreter( model_path=self.model_path, experimental_delegates=[ tf.lite.experimental.load_delegate(self.EDGETPU_SHARED_LIB) - ] + ], ) self.tflite_interpreter.allocate_tensors() @@ -71,28 +75,32 @@ def __init__( self.output_details = self.tflite_interpreter.get_output_details() self.category_index = create_category_index_from_labelmap( - self.PATH_TO_LABELS, use_display_name=True) + self.PATH_TO_LABELS, use_display_name=True + ) logging.info( - f'loaded labels from {self.PATH_TO_LABELS} \n {self.category_index}') + f"loaded labels from {self.PATH_TO_LABELS} \n {self.category_index}" + ) - logging.info(f'initialized model {model_name} \n') - logging.info( - f'model inputs: {self.input_details} \n {self.input_details}') - logging.info( - f'model outputs: {self.output_details} \n {self.output_details}') + logging.info(f"initialized model {model_name} \n") + logging.info(f"model inputs: {self.input_details} \n {self.input_details}") + logging.info(f"model outputs: {self.output_details} \n {self.output_details}") + + @classmethod + def validate_labels(cls, labels): + return all([x in cls.LABELS for x in labels]) def label_to_category_index(self, labels): # @todo :trashfire: - return tuple(map( - lambda x: x['id'], - filter( - lambda x: x['name'] in labels, self.category_index.values() + return tuple( + map( + lambda x: x["id"], + filter(lambda x: x["name"] in labels, self.category_index.values()), ) - )) + ) def label_display_name_by_idx(self, idx): - return self.category_index[idx]['display_name'] + return self.category_index[idx]["display_name"] def create_overlay(self, image_np, output_dict): @@ -101,14 +109,14 @@ def create_overlay(self, image_np, output_dict): # draw bounding boxes visualize_boxes_and_labels_on_image_array( image_np, - output_dict['detection_boxes'], - output_dict['detection_classes'], - output_dict['detection_scores'], + output_dict["detection_boxes"], + output_dict["detection_classes"], + output_dict["detection_scores"], self.category_index, use_normalized_coordinates=True, line_thickness=4, min_score_thresh=self.min_score_thresh, - max_boxes_to_draw=3 + max_boxes_to_draw=3, ) img = Image.fromarray(image_np) @@ -116,20 +124,20 @@ def create_overlay(self, image_np, output_dict): return img.tobytes() def predict(self, image): - ''' - image - np.array (3 RGB channels) + """ + image - np.array (3 RGB channels) - returns - { - 'detection_classes': int64, - 'num_detections': int64 - 'detection_masks': ... - } - ''' + returns + { + 'detection_classes': int64, + 'num_detections': int64 + 'detection_masks': ... + } + """ image = np.asarray(image) # normalize 0 - 255 RGB to values between (-1, 1) - #image = (image / 128.0) - 1 + # image = (image / 128.0) - 1 # The input needs to be a tensor, convert it using `tf.convert_to_tensor`. @@ -139,8 +147,7 @@ def predict(self, image): input_tensor = input_tensor[tf.newaxis, ...] # Run inference - self.tflite_interpreter.set_tensor( - self.input_details[0]['index'], input_tensor) + self.tflite_interpreter.set_tensor(self.input_details[0]["index"], input_tensor) self.tflite_interpreter.invoke() @@ -161,44 +168,47 @@ def predict(self, image): # [1, num_anchors, num_classes] containing the class scores for each anchor # after applying score conversion. - box_data = tf.convert_to_tensor(self.tflite_interpreter.get_tensor( - self.output_details[0]['index'])) - class_data = tf.convert_to_tensor(self.tflite_interpreter.get_tensor( - self.output_details[1]['index'])) - score_data = tf.convert_to_tensor(self.tflite_interpreter.get_tensor( - self.output_details[2]['index'])) - num_detections = tf.convert_to_tensor(self.tflite_interpreter.get_tensor( - self.output_details[3]['index'])) - - class_data = tf.squeeze( - class_data, axis=[0]).numpy().astype(np.int64) + 1 + box_data = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[0]["index"]) + ) + class_data = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[1]["index"]) + ) + score_data = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[2]["index"]) + ) + num_detections = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[3]["index"]) + ) + + class_data = tf.squeeze(class_data, axis=[0]).numpy().astype(np.int64) + 1 box_data = tf.squeeze(box_data, axis=[0]).numpy() score_data = tf.squeeze(score_data, axis=[0]).numpy() return { - 'detection_boxes': box_data, - 'detection_classes': class_data, - 'detection_scores': score_data, - 'num_detections': len(num_detections) + "detection_boxes": box_data, + "detection_classes": class_data, + "detection_scores": score_data, + "num_detections": len(num_detections), } -class FaceSSD_MobileNet_V2(object): +class FaceSSDMobileNetV2Float32(object): - PATH_TO_LABELS = rpi_deep_pantilt_path[0] + '/data/facessd_label_map.pbtxt' + PATH_TO_LABELS = rpi_deep_pantilt_path[0] + "/data/facessd_label_map.pbtxt" + LABELS = ["face"] def __init__( self, - base_url='https://github.com/leigh-johnson/rpi-deep-pantilt/releases/download/v1.0.1/', - model_name='facessd_mobilenet_v2_quantized_320x320_open_image_v4_tflite2', + base_url="https://github.com/leigh-johnson/rpi-deep-pantilt/releases/download/v1.0.1/", + model_name="facessd_mobilenet_v2_quantized_320x320_open_image_v4_tflite2", input_shape=(320, 320), - min_score_thresh=0.6 - + min_score_thresh=0.6, ): self.base_url = base_url self.model_name = model_name - self.model_file = model_name + '.tar.gz' + self.model_file = model_name + ".tar.gz" self.model_url = base_url + self.model_file self.min_score_thresh = min_score_thresh @@ -206,12 +216,13 @@ def __init__( fname=self.model_name, origin=self.model_url, untar=True, - cache_subdir='models' + cache_subdir="models", ) - self.model_path = os.path.splitext( - os.path.splitext(self.model_dir)[0] - )[0] + '/model_postprocessed.tflite' + self.model_path = ( + os.path.splitext(os.path.splitext(self.model_dir)[0])[0] + + "/model_postprocessed.tflite" + ) self.tflite_interpreter = tf.lite.Interpreter( model_path=self.model_path, @@ -222,28 +233,32 @@ def __init__( self.output_details = self.tflite_interpreter.get_output_details() self.category_index = create_category_index_from_labelmap( - self.PATH_TO_LABELS, use_display_name=True) + self.PATH_TO_LABELS, use_display_name=True + ) logging.info( - f'loaded labels from {self.PATH_TO_LABELS} \n {self.category_index}') + f"loaded labels from {self.PATH_TO_LABELS} \n {self.category_index}" + ) - logging.info(f'initialized model {model_name} \n') - logging.info( - f'model inputs: {self.input_details} \n {self.input_details}') - logging.info( - f'model outputs: {self.output_details} \n {self.output_details}') + logging.info(f"initialized model {model_name} \n") + logging.info(f"model inputs: {self.input_details} \n {self.input_details}") + logging.info(f"model outputs: {self.output_details} \n {self.output_details}") + + @classmethod + def validate_labels(cls, labels): + return all([x in cls.LABELS for x in labels]) def label_to_category_index(self, labels): # @todo :trashfire: - return tuple(map( - lambda x: x['id'], - filter( - lambda x: x['name'] in labels, self.category_index.values() + return tuple( + map( + lambda x: x["id"], + filter(lambda x: x["name"] in labels, self.category_index.values()), ) - )) + ) def label_display_name_by_idx(self, idx): - return self.category_index[idx]['display_name'] + return self.category_index[idx]["display_name"] def create_overlay(self, image_np, output_dict): @@ -252,14 +267,14 @@ def create_overlay(self, image_np, output_dict): # draw bounding boxes visualize_boxes_and_labels_on_image_array( image_np, - output_dict['detection_boxes'], - output_dict['detection_classes'], - output_dict['detection_scores'], + output_dict["detection_boxes"], + output_dict["detection_classes"], + output_dict["detection_scores"], self.category_index, use_normalized_coordinates=True, line_thickness=4, min_score_thresh=self.min_score_thresh, - max_boxes_to_draw=3 + max_boxes_to_draw=3, ) img = Image.fromarray(image_np) @@ -267,31 +282,29 @@ def create_overlay(self, image_np, output_dict): return img.tobytes() def predict(self, image): - ''' - image - np.array (3 RGB channels) + """ + image - np.array (3 RGB channels) - returns - { - 'detection_classes': int64, - 'num_detections': int64 - 'detection_masks': ... - } - ''' + returns + { + 'detection_classes': int64, + 'num_detections': int64 + 'detection_masks': ... + } + """ image = np.asarray(image) # normalize 0 - 255 RGB to values between (-1, 1) image = (image / 128.0) - 1 # The input needs to be a tensor, convert it using `tf.convert_to_tensor`. - input_tensor = tf.convert_to_tensor(image, dtype=tf.float32) # The model expects a batch of images, so add an axis with `tf.newaxis`. input_tensor = input_tensor[tf.newaxis, ...] # Run inference - self.tflite_interpreter.set_tensor( - self.input_details[0]['index'], input_tensor) + self.tflite_interpreter.set_tensor(self.input_details[0]["index"], input_tensor) self.tflite_interpreter.invoke() @@ -312,24 +325,27 @@ def predict(self, image): # [1, num_anchors, num_classes] containing the class scores for each anchor # after applying score conversion. - box_data = tf.convert_to_tensor(self.tflite_interpreter.get_tensor( - self.output_details[0]['index'])) - class_data = tf.convert_to_tensor(self.tflite_interpreter.get_tensor( - self.output_details[1]['index'])) - score_data = tf.convert_to_tensor(self.tflite_interpreter.get_tensor( - self.output_details[2]['index'])) - num_detections = tf.convert_to_tensor(self.tflite_interpreter.get_tensor( - self.output_details[3]['index'])) + box_data = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[0]["index"]) + ) + class_data = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[1]["index"]) + ) + score_data = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[2]["index"]) + ) + num_detections = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[3]["index"]) + ) # hilarious, but it seems like all classes predictions are off by 1 idx - class_data = tf.squeeze( - class_data, axis=[0]).numpy().astype(np.int64) + 1 + class_data = tf.squeeze(class_data, axis=[0]).numpy().astype(np.int64) + 1 box_data = tf.squeeze(box_data, axis=[0]).numpy() score_data = tf.squeeze(score_data, axis=[0]).numpy() return { - 'detection_boxes': box_data, # 10, 4 - 'detection_classes': class_data, # 10 - 'detection_scores': score_data, # 10, - 'num_detections': len(num_detections) + "detection_boxes": box_data, # 10, 4 + "detection_classes": class_data, # 10 + "detection_scores": score_data, # 10, + "num_detections": len(num_detections), } diff --git a/rpi_deep_pantilt/detect/pretrained/api_v1/ssd_mobilenet_v3_coco.py b/rpi_deep_pantilt/detect/pretrained/api_v1/ssd_mobilenet_v3_coco.py new file mode 100644 index 0000000..0149242 --- /dev/null +++ b/rpi_deep_pantilt/detect/pretrained/api_v1/ssd_mobilenet_v3_coco.py @@ -0,0 +1,543 @@ +# Python +import logging +import pathlib +import os +import sys + +# lib +import numpy as np +from PIL import Image +import tensorflow as tf + +from rpi_deep_pantilt import __path__ as rpi_deep_pantilt_path +from rpi_deep_pantilt.detect.util.exceptions import InvalidLabelException +from rpi_deep_pantilt.detect.util.label import create_category_index_from_labelmap +from rpi_deep_pantilt.detect.util.visualization import ( + visualize_boxes_and_labels_on_image_array, +) + + +class SSDMobileNetV3EdgeTPU(object): + + EDGETPU_SHARED_LIB = "libedgetpu.so.1" + PATH_TO_LABELS = rpi_deep_pantilt_path[0] + "/data/mscoco_label_map.pbtxt" + LABELS = [ + "person", + "bicycle", + "car", + "motorcycle", + "airplane", + "bus", + "train", + "truck", + "boat", + "traffic light", + "fire hydrant", + "stop sign", + "parking meter", + "bench", + "bird", + "cat", + "dog", + "horse", + "sheep", + "cow", + "elephant", + "bear", + "zebra", + "giraffe", + "backpack", + "umbrella", + "handbag", + "tie", + "suitcase", + "frisbee", + "skis", + "snowboard", + "sports ball", + "kite", + "baseball bat", + "baseball glove", + "skateboard", + "surfboard", + "tennis racket", + "bottle", + "wine glass", + "cup", + "fork", + "knife", + "spoon", + "bowl", + "banana", + "apple", + "sandwich", + "orange", + "broccoli", + "carrot", + "hot dog", + "pizza", + "donut", + "cake", + "chair", + "couch", + "potted plant", + "bed", + "dining table", + "toilet", + "tv", + "laptop", + "mouse", + "remote", + "keyboard", + "cell phone", + "microwave", + "oven", + "toaster", + "sink", + "refrigerator", + "book", + "clock", + "vase", + "scissors", + "teddy bear", + "hair drier", + "toothbrush", + ] + + def __init__( + self, + base_url="https://github.com/leigh-johnson/rpi-deep-pantilt/releases/download/v1.0.0/", + model_name="ssdlite_mobilenet_edgetpu_coco_quant", + input_shape=(320, 320), + min_score_thresh=0.50, + tflite_model_file="model_postprocessed_quantized_128_uint8_edgetpu.tflite", + ): + + self.base_url = base_url + self.model_name = model_name + self.model_file = model_name + ".tar.gz" + self.model_url = base_url + self.model_file + self.tflite_model_file = tflite_model_file + + self.model_dir = tf.keras.utils.get_file( + fname=self.model_file, + origin=self.model_url, + untar=True, + cache_subdir="models", + ) + + self.min_score_thresh = min_score_thresh + + self.model_path = ( + os.path.splitext(os.path.splitext(self.model_dir)[0])[0] + + f"/{self.tflite_model_file}" + ) + + try: + from tflite_runtime import interpreter as coral_tflite_interpreter + except ImportError as e: + logging.error(e) + logging.error("Please install Edge TPU tflite_runtime:") + logging.error( + "$ pip install https://dl.google.com/coral/python/tflite_runtime-2.1.0.post1-cp37-cp37m-linux_armv7l.whl" + ) + sys.exit(1) + + self.tflite_interpreter = coral_tflite_interpreter.Interpreter( + model_path=self.model_path, + experimental_delegates=[ + tf.lite.experimental.load_delegate(self.EDGETPU_SHARED_LIB) + ], + ) + + self.tflite_interpreter.allocate_tensors() + + self.input_details = self.tflite_interpreter.get_input_details() + self.output_details = self.tflite_interpreter.get_output_details() + + self.category_index = create_category_index_from_labelmap( + self.PATH_TO_LABELS, use_display_name=True + ) + + logging.info( + f"loaded labels from {self.PATH_TO_LABELS} \n {self.category_index}" + ) + + logging.info(f"initialized model {model_name} \n") + logging.info(f"model inputs: {self.input_details} \n {self.input_details}") + logging.info(f"model outputs: {self.output_details} \n {self.output_details}") + + @classmethod + def validate_labels(cls, labels): + return all([x in cls.LABELS for x in labels]) + + def label_to_category_index(self, labels): + return tuple( + map( + lambda x: x["id"], + filter(lambda x: x["name"] in labels, self.category_index.values()), + ) + ) + + def label_display_name_by_idx(self, idx): + return self.category_index[idx]["display_name"] + + def filter_tracked(self, prediction, label_idxs): + """ + Zero predictions not in list of label_idxs + """ + return { + "detection_boxes": prediction.get("detection_boxes"), + "detection_classes": prediction.get("detection_classes"), + "detection_scores": np.array( + [ + v if prediction.get("detection_classes")[i] in label_idxs else 0.0 + for i, v in enumerate(prediction.get("detection_scores")) + ] + ), + } + + def create_overlay(self, image_np, output_dict): + + image_np = image_np.copy() + + # draw bounding boxes + visualize_boxes_and_labels_on_image_array( + image_np, + output_dict["detection_boxes"], + output_dict["detection_classes"], + output_dict["detection_scores"], + self.category_index, + use_normalized_coordinates=True, + line_thickness=4, + min_score_thresh=self.min_score_thresh, + max_boxes_to_draw=3, + ) + + img = Image.fromarray(image_np) + + return img.tobytes() + + def predict(self, image): + """ + image - np.array (3 RGB channels) + + returns + { + 'detection_classes': int64, + 'num_detections': int64 + 'detection_masks': ... + } + """ + + image = np.asarray(image) + # normalize 0 - 255 RGB to values between (-1, 1) + # image = (image / 128.0) - 1 + + # The input needs to be a tensor, convert it using `tf.convert_to_tensor`. + + input_tensor = tf.convert_to_tensor(image, dtype=tf.uint8) + + # The model expects a batch of images, so add an axis with `tf.newaxis`. + input_tensor = input_tensor[tf.newaxis, ...] + + # Run inference + self.tflite_interpreter.set_tensor(self.input_details[0]["index"], input_tensor) + + self.tflite_interpreter.invoke() + + # TFLite_Detection_PostProcess custom op node has four outputs: + # detection_boxes: a float32 tensor of shape [1, num_boxes, 4] with box + # locations + # detection_classes: a float32 tensor of shape [1, num_boxes] + # with class indices + # detection_scores: a float32 tensor of shape [1, num_boxes] + # with class scores + # num_boxes: a float32 tensor of size 1 containing the number of detected + # boxes + + # Without the PostProcessing ops, the graph has two outputs: + # 'raw_outputs/box_encodings': a float32 tensor of shape [1, num_anchors, 4] + # containing the encoded box predictions. + # 'raw_outputs/class_predictions': a float32 tensor of shape + # [1, num_anchors, num_classes] containing the class scores for each anchor + # after applying score conversion. + + box_data = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[0]["index"]) + ) + class_data = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[1]["index"]) + ) + score_data = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[2]["index"]) + ) + num_detections = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[3]["index"]) + ) + + class_data = tf.squeeze(class_data, axis=[0]).numpy().astype(np.int64) + 1 + box_data = tf.squeeze(box_data, axis=[0]).numpy() + score_data = tf.squeeze(score_data, axis=[0]).numpy() + + return { + "detection_boxes": box_data, + "detection_classes": class_data, + "detection_scores": score_data, + "num_detections": len(num_detections), + } + + +class SSDMobileNetV3Float32(object): + + PATH_TO_LABELS = rpi_deep_pantilt_path[0] + "/data/mscoco_label_map.pbtxt" + LABELS = [ + "person", + "bicycle", + "car", + "motorcycle", + "airplane", + "bus", + "train", + "truck", + "boat", + "traffic light", + "fire hydrant", + "stop sign", + "parking meter", + "bench", + "bird", + "cat", + "dog", + "horse", + "sheep", + "cow", + "elephant", + "bear", + "zebra", + "giraffe", + "backpack", + "umbrella", + "handbag", + "tie", + "suitcase", + "frisbee", + "skis", + "snowboard", + "sports ball", + "kite", + "baseball bat", + "baseball glove", + "skateboard", + "surfboard", + "tennis racket", + "bottle", + "wine glass", + "cup", + "fork", + "knife", + "spoon", + "bowl", + "banana", + "apple", + "sandwich", + "orange", + "broccoli", + "carrot", + "hot dog", + "pizza", + "donut", + "cake", + "chair", + "couch", + "potted plant", + "bed", + "dining table", + "toilet", + "tv", + "laptop", + "mouse", + "remote", + "keyboard", + "cell phone", + "microwave", + "oven", + "toaster", + "sink", + "refrigerator", + "book", + "clock", + "vase", + "scissors", + "teddy bear", + "hair drier", + "toothbrush", + ] + + def __init__( + self, + base_url="https://github.com/leigh-johnson/rpi-deep-pantilt/releases/download/v1.0.0/", + model_name="ssd_mobilenet_v3_small_coco_2019_08_14", + input_shape=(320, 320), + min_score_thresh=0.6, + ): + + self.base_url = base_url + self.model_name = model_name + self.model_file = model_name + ".tar.gz" + self.model_url = base_url + self.model_file + self.min_score_thresh = min_score_thresh + + self.model_dir = tf.keras.utils.get_file( + fname=self.model_name, + origin=self.model_url, + untar=True, + cache_subdir="models", + ) + + self.model_path = ( + os.path.splitext(os.path.splitext(self.model_dir)[0])[0] + + "/model_postprocessed.tflite" + ) + + self.tflite_interpreter = tf.lite.Interpreter( + model_path=self.model_path, + ) + self.tflite_interpreter.allocate_tensors() + + self.input_details = self.tflite_interpreter.get_input_details() + self.output_details = self.tflite_interpreter.get_output_details() + + self.category_index = create_category_index_from_labelmap( + self.PATH_TO_LABELS, use_display_name=True + ) + + logging.info( + f"loaded labels from {self.PATH_TO_LABELS} \n {self.category_index}" + ) + + logging.info(f"initialized model {model_name} \n") + logging.info(f"model inputs: {self.input_details} \n {self.input_details}") + logging.info(f"model outputs: {self.output_details} \n {self.output_details}") + + @classmethod + def validate_labels(cls, labels): + return all([x in cls.LABELS for x in labels]) + + def label_to_category_index(self, labels): + # @todo :trashfire: + return tuple( + map( + lambda x: x["id"], + filter(lambda x: x["name"] in labels, self.category_index.values()), + ) + ) + + def label_display_name_by_idx(self, idx): + return self.category_index[idx]["display_name"] + + def filter_tracked(self, prediction, label_idxs): + """ + Zero predictions not in list of label_idxs + """ + return { + "detection_boxes": prediction.get("detection_boxes"), + "detection_classes": prediction.get("detection_classes"), + "detection_scores": np.array( + [ + v if prediction.get("detection_classes")[i] in label_idxs else 0.0 + for i, v in enumerate(prediction.get("detection_scores")) + ] + ), + } + + def create_overlay(self, image_np, output_dict): + + image_np = image_np.copy() + + # draw bounding boxes + visualize_boxes_and_labels_on_image_array( + image_np, + output_dict["detection_boxes"], + output_dict["detection_classes"], + output_dict["detection_scores"], + self.category_index, + use_normalized_coordinates=True, + line_thickness=4, + min_score_thresh=self.min_score_thresh, + max_boxes_to_draw=3, + ) + + img = Image.fromarray(image_np) + + return img.tobytes() + + def predict(self, image): + """ + image - np.array (3 RGB channels) + + returns + { + 'detection_classes': int64, + 'num_detections': int64 + 'detection_masks': ... + } + """ + + image = np.asarray(image) + # normalize 0 - 255 RGB to values between (-1, 1) + image = (image / 128.0) - 1 + + # The input needs to be a tensor, convert it using `tf.convert_to_tensor`. + + input_tensor = tf.convert_to_tensor(image, dtype=tf.float32) + + # The model expects a batch of images, so add an axis with `tf.newaxis`. + input_tensor = input_tensor[tf.newaxis, ...] + + # Run inference + self.tflite_interpreter.set_tensor(self.input_details[0]["index"], input_tensor) + + self.tflite_interpreter.invoke() + + # TFLite_Detection_PostProcess custom op node has four outputs: + # detection_boxes: a float32 tensor of shape [1, num_boxes, 4] with box + # locations + # detection_classes: a float32 tensor of shape [1, num_boxes] + # with class indices + # detection_scores: a float32 tensor of shape [1, num_boxes] + # with class scores + # num_boxes: a float32 tensor of size 1 containing the number of detected + # boxes + + # Without the PostProcessing ops, the graph has two outputs: + # 'raw_outputs/box_encodings': a float32 tensor of shape [1, num_anchors, 4] + # containing the encoded box predictions. + # 'raw_outputs/class_predictions': a float32 tensor of shape + # [1, num_anchors, num_classes] containing the class scores for each anchor + # after applying score conversion. + + box_data = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[0]["index"]) + ) + class_data = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[1]["index"]) + ) + score_data = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[2]["index"]) + ) + num_detections = tf.convert_to_tensor( + self.tflite_interpreter.get_tensor(self.output_details[3]["index"]) + ) + + # hilarious, but it seems like all classes predictions are off by 1 idx + class_data = tf.squeeze(class_data, axis=[0]).numpy().astype(np.int64) + 1 + box_data = tf.squeeze(box_data, axis=[0]).numpy() + score_data = tf.squeeze(score_data, axis=[0]).numpy() + + return { + "detection_boxes": box_data, # 10, 4 + "detection_classes": class_data, # 10 + "detection_scores": score_data, # 10, + "num_detections": len(num_detections), + } diff --git a/rpi_deep_pantilt/detect/pretrained/api_v2/__init__.py b/rpi_deep_pantilt/detect/pretrained/api_v2/__init__.py new file mode 100644 index 0000000..481be5d --- /dev/null +++ b/rpi_deep_pantilt/detect/pretrained/api_v2/__init__.py @@ -0,0 +1,15 @@ +from rpi_deep_pantilt.detect.pretrained.api_v2.facessd_mobilenet_v2 import ( + FaceSSDMobileNetV2EdgeTPU, + FaceSSDMobileNetV2Float32, + FaceSSDMobileNetV2Int8, +) + +from rpi_deep_pantilt.detect.pretrained.api_v2.ssd_mobilenet_v3_coco import ( + SSDMobileNetV3EdgeTPU, + SSDMobileNetV3Float32, + SSDMobileNetV3Int8, +) + +from rpi_deep_pantilt.detect.pretrained.api_v2.automl_leopard_example import ( + LeopardAutoMLInt8, +) diff --git a/rpi_deep_pantilt/detect/pretrained/api_v2/automl_leopard_example.py b/rpi_deep_pantilt/detect/pretrained/api_v2/automl_leopard_example.py new file mode 100644 index 0000000..8725358 --- /dev/null +++ b/rpi_deep_pantilt/detect/pretrained/api_v2/automl_leopard_example.py @@ -0,0 +1,36 @@ +import tensorflow as tf + +from rpi_deep_pantilt import __path__ as rpi_deep_pantilt_path +from rpi_deep_pantilt.detect.custom.base_predictors import ( + TFLiteDetectionPostProcessPredictor, +) + + +class LeopardAutoMLInt8(TFLiteDetectionPostProcessPredictor): + + LABELS = ["leopard"] + + def __init__( + self, + model_name="leopard_detector", + tflite_file=rpi_deep_pantilt_path[0] + + "/data/leopard.tflite", # substitute the FULL path to your .tflite file + label_file=rpi_deep_pantilt_path[0] + + "/data/leopard_label_map.pbtxt", # you only need to change this if adding additional labels (background class is implicitly id:0) + input_shape=( + 320, + 320, + ), # make sure height, width match metadata.json inputShape + min_score_thresh=0.95, # increase threshold to reduce false positive results + input_type=tf.uint8, # do not change + ): + + super().__init__( + model_name=model_name, + tflite_file=tflite_file, + label_file=label_file, + model_uri=None, # model_uri is used to download remote files; None implies tflite_File and label_file are fully qualified local paths + input_shape=input_shape, + min_score_thresh=min_score_thresh, + input_type=input_type, + ) diff --git a/rpi_deep_pantilt/detect/pretrained/api_v2/facessd_mobilenet_v2.py b/rpi_deep_pantilt/detect/pretrained/api_v2/facessd_mobilenet_v2.py new file mode 100644 index 0000000..5d9efdc --- /dev/null +++ b/rpi_deep_pantilt/detect/pretrained/api_v2/facessd_mobilenet_v2.py @@ -0,0 +1,94 @@ +import tensorflow as tf + +from rpi_deep_pantilt import __path__ as rpi_deep_pantilt_path +from rpi_deep_pantilt.detect.custom.base_predictors import ( + TFLiteDetectionPostProcessPredictor, +) + + +class FaceSSDMobileNetV2EdgeTPU(TFLiteDetectionPostProcessPredictor): + """ + Model source: https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf1_detection_zoo.md#open-images-trained-models + + Non-max supression op (TFLite_Detection_Postprocess) added to graph via tools/tflite-postprocess-ops-128-uint8-quant.sh + """ + + LABELS = ["face"] + + def __init__( + self, + model_uri="https://github.com/leigh-johnson/rpi-deep-pantilt/releases/download/v1.1.1/facessd_mobilenet_v2_quantized_320x320_open_image_v4_tflite2.tar.gz", + model_name="facessd_mobilenet_v2_quantized_320x320_open_image_v4_tflite2", + input_shape=(320, 320), + min_score_thresh=0.50, + input_type=tf.uint8, + tflite_file="model_postprocessed_quantized_128_uint8_edgetpu.tflite", + label_file=rpi_deep_pantilt_path[0] + "/data/facessd_label_map.pbtxt", + ): + + super().__init__( + model_name=model_name, + tflite_file=tflite_file, + label_file=label_file, + model_uri=model_uri, + input_shape=input_shape, + min_score_thresh=min_score_thresh, + input_type=input_type, + edge_tpu=True, + ) + + +class FaceSSDMobileNetV2Int8(TFLiteDetectionPostProcessPredictor): + """ + Model source: https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf1_detection_zoo.md#open-images-trained-models + + Non-max supression op (TFLite_Detection_Postprocess) added to graph via tools/tflite-postprocess-ops-128-uint8-quant.sh + """ + + LABELS = ["face"] + + def __init__( + self, + model_uri="https://github.com/leigh-johnson/rpi-deep-pantilt/releases/download/v1.1.1/facessd_mobilenet_v2_quantized_320x320_open_image_v4_tflite2.tar.gz", + model_name="facessd_mobilenet_v2_quantized_320x320_open_image_v4_tflite2", + input_shape=(320, 320), + min_score_thresh=0.50, + input_type=tf.uint8, + tflite_file="model_postprocessed_quantized_128_uint8.tflite", + label_file=rpi_deep_pantilt_path[0] + "/data/facessd_label_map.pbtxt", + ): + + super().__init__( + model_name=model_name, + tflite_file=tflite_file, + label_file=label_file, + model_uri=model_uri, + input_shape=input_shape, + min_score_thresh=min_score_thresh, + input_type=input_type, + ) + + +class FaceSSDMobileNetV2Float32(TFLiteDetectionPostProcessPredictor): + LABELS = ["face"] + + def __init__( + self, + model_uri="https://github.com/leigh-johnson/rpi-deep-pantilt/releases/download/v1.1.1/facessd_mobilenet_v2_quantized_320x320_open_image_v4_tflite2.tar.gz", + model_name="facessd_mobilenet_v2_quantized_320x320_open_image_v4_tflite2", + input_shape=(320, 320), + min_score_thresh=0.50, + input_type=tf.float32, + tflite_file="model_postprocessed.tflite", + label_file=rpi_deep_pantilt_path[0] + "/data/facessd_label_map.pbtxt", + ): + + super().__init__( + model_name=model_name, + tflite_file=tflite_file, + label_file=label_file, + model_uri=model_uri, + input_shape=input_shape, + min_score_thresh=min_score_thresh, + input_type=input_type, + ) diff --git a/rpi_deep_pantilt/detect/pretrained/api_v2/ssd_mobilenet_v3_coco.py b/rpi_deep_pantilt/detect/pretrained/api_v2/ssd_mobilenet_v3_coco.py new file mode 100644 index 0000000..cdd3690 --- /dev/null +++ b/rpi_deep_pantilt/detect/pretrained/api_v2/ssd_mobilenet_v3_coco.py @@ -0,0 +1,337 @@ +import tensorflow as tf + +from rpi_deep_pantilt import __path__ as rpi_deep_pantilt_path +from rpi_deep_pantilt.detect.custom.base_predictors import ( + TFLiteDetectionPostProcessPredictor, +) + + +class SSDMobileNetV3EdgeTPU(TFLiteDetectionPostProcessPredictor): + """ + Model source: https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf1_detection_zoo.md#open-images-trained-models + + Non-max supression op (TFLite_Detection_Postprocess) added to graph via tools/tflite-postprocess-ops-128-uint8-quant.sh + """ + + LABELS = [ + "person", + "bicycle", + "car", + "motorcycle", + "airplane", + "bus", + "train", + "truck", + "boat", + "traffic light", + "fire hydrant", + "stop sign", + "parking meter", + "bench", + "bird", + "cat", + "dog", + "horse", + "sheep", + "cow", + "elephant", + "bear", + "zebra", + "giraffe", + "backpack", + "umbrella", + "handbag", + "tie", + "suitcase", + "frisbee", + "skis", + "snowboard", + "sports ball", + "kite", + "baseball bat", + "baseball glove", + "skateboard", + "surfboard", + "tennis racket", + "bottle", + "wine glass", + "cup", + "fork", + "knife", + "spoon", + "bowl", + "banana", + "apple", + "sandwich", + "orange", + "broccoli", + "carrot", + "hot dog", + "pizza", + "donut", + "cake", + "chair", + "couch", + "potted plant", + "bed", + "dining table", + "toilet", + "tv", + "laptop", + "mouse", + "remote", + "keyboard", + "cell phone", + "microwave", + "oven", + "toaster", + "sink", + "refrigerator", + "book", + "clock", + "vase", + "scissors", + "teddy bear", + "hair drier", + "toothbrush", + ] + + def __init__( + self, + model_uri="https://github.com/leigh-johnson/rpi-deep-pantilt/releases/download/v1.1.1/ssdlite_mobilenet_edgetpu_coco_quant.tar.gz", + model_name="ssdlite_mobilenet_edgetpu_coco_quant", + input_shape=(320, 320), + min_score_thresh=0.50, + input_type=tf.uint8, + tflite_file="model_postprocessed_quantized_128_uint8_edgetpu.tflite", + label_file=rpi_deep_pantilt_path[0] + "/data/mscoco_label_map.pbtxt", + ): + + super().__init__( + model_name=model_name, + tflite_file=tflite_file, + label_file=label_file, + model_uri=model_uri, + input_shape=input_shape, + min_score_thresh=min_score_thresh, + input_type=input_type, + edge_tpu=True, + ) + + +class SSDMobileNetV3Int8(TFLiteDetectionPostProcessPredictor): + """ + Model source: https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf1_detection_zoo.md#open-images-trained-models + + Non-max supression op (TFLite_Detection_Postprocess) added to graph via tools/tflite-postprocess-ops-128-uint8-quant.sh + """ + + LABELS = [ + "person", + "bicycle", + "car", + "motorcycle", + "airplane", + "bus", + "train", + "truck", + "boat", + "traffic light", + "fire hydrant", + "stop sign", + "parking meter", + "bench", + "bird", + "cat", + "dog", + "horse", + "sheep", + "cow", + "elephant", + "bear", + "zebra", + "giraffe", + "backpack", + "umbrella", + "handbag", + "tie", + "suitcase", + "frisbee", + "skis", + "snowboard", + "sports ball", + "kite", + "baseball bat", + "baseball glove", + "skateboard", + "surfboard", + "tennis racket", + "bottle", + "wine glass", + "cup", + "fork", + "knife", + "spoon", + "bowl", + "banana", + "apple", + "sandwich", + "orange", + "broccoli", + "carrot", + "hot dog", + "pizza", + "donut", + "cake", + "chair", + "couch", + "potted plant", + "bed", + "dining table", + "toilet", + "tv", + "laptop", + "mouse", + "remote", + "keyboard", + "cell phone", + "microwave", + "oven", + "toaster", + "sink", + "refrigerator", + "book", + "clock", + "vase", + "scissors", + "teddy bear", + "hair drier", + "toothbrush", + ] + + def __init__( + self, + model_uri="https://github.com/leigh-johnson/rpi-deep-pantilt/releases/download/v1.1.1/ssdlite_mobilenet_edgetpu_coco_quant.tar.gz", + model_name="ssd_mobilenet_v3_small_coco_2019_08_14", + input_shape=(320, 320), + min_score_thresh=0.50, + input_type=tf.uint8, + tflite_file="model_postprocessed_quantized_128_uint8.tflite", + label_file=rpi_deep_pantilt_path[0] + "/data/mscoco_label_map.pbtxt", + ): + + super().__init__( + model_name=model_name, + tflite_file=tflite_file, + label_file=label_file, + model_uri=model_uri, + input_shape=input_shape, + min_score_thresh=min_score_thresh, + input_type=input_type, + ) + + +class SSDMobileNetV3Float32(TFLiteDetectionPostProcessPredictor): + LABELS = [ + "person", + "bicycle", + "car", + "motorcycle", + "airplane", + "bus", + "train", + "truck", + "boat", + "traffic light", + "fire hydrant", + "stop sign", + "parking meter", + "bench", + "bird", + "cat", + "dog", + "horse", + "sheep", + "cow", + "elephant", + "bear", + "zebra", + "giraffe", + "backpack", + "umbrella", + "handbag", + "tie", + "suitcase", + "frisbee", + "skis", + "snowboard", + "sports ball", + "kite", + "baseball bat", + "baseball glove", + "skateboard", + "surfboard", + "tennis racket", + "bottle", + "wine glass", + "cup", + "fork", + "knife", + "spoon", + "bowl", + "banana", + "apple", + "sandwich", + "orange", + "broccoli", + "carrot", + "hot dog", + "pizza", + "donut", + "cake", + "chair", + "couch", + "potted plant", + "bed", + "dining table", + "toilet", + "tv", + "laptop", + "mouse", + "remote", + "keyboard", + "cell phone", + "microwave", + "oven", + "toaster", + "sink", + "refrigerator", + "book", + "clock", + "vase", + "scissors", + "teddy bear", + "hair drier", + "toothbrush", + ] + + def __init__( + self, + model_uri="https://github.com/leigh-johnson/rpi-deep-pantilt/releases/download/v1.1.1/ssd_mobilenet_v3_small_coco_2019_08_14.tar.gz", + model_name="ssd_mobilenet_v3_small_coco_2019_08_14", + input_shape=(320, 320), + min_score_thresh=0.50, + input_type=tf.float32, + tflite_file="model_postprocessed_quantized.tflite", + label_file=rpi_deep_pantilt_path[0] + "/data/mscoco_label_map.pbtxt", + ): + + super().__init__( + model_name=model_name, + tflite_file=tflite_file, + label_file=label_file, + model_uri=model_uri, + input_shape=input_shape, + min_score_thresh=min_score_thresh, + input_type=input_type, + ) diff --git a/rpi_deep_pantilt/detect/registry.py b/rpi_deep_pantilt/detect/registry.py new file mode 100644 index 0000000..8a610e4 --- /dev/null +++ b/rpi_deep_pantilt/detect/registry.py @@ -0,0 +1,67 @@ +import importlib +import logging + +from rpi_deep_pantilt.detect.util.exceptions import InvalidLabelException + + +class ModelRegistry(object): + + FLOAT32_CLASSES = ( + "FaceSSDMobileNetV2Float32", + "SSDMobileNetV3Float32", + ) + + UINT8_CLASSES = ( + "FaceSSDMobileNetV2Int8", + "SSDMobileNetV3Int8", + "LeopardAutoMLInt8", + ) + + EDGETPU_CLASSES = ( + "SSDMobileNetV3EdgeTPU", + "FaceSSDMobileNetV2EdgeTPU", + ) + + def __init__(self, edge_tpu, api_version, dtype): + self.edge_tpu = edge_tpu + self.api_version = api_version + self.version_str = f"api_{api_version}" + + self.import_path = f"rpi_deep_pantilt.detect.pretrained.{self.version_str}" + + self.module = importlib.import_module(self.import_path) + + if edge_tpu: + self.model_list = self.EDGETPU_CLASSES + self.default_model = self.module.SSDMobileNetV3EdgeTPU + elif dtype == "uint8": + self.model_list = self.UINT8_CLASSES + self.default_model = self.module.SSDMobileNetV3Int8 + else: + self.model_list = self.FLOAT32_CLASSES + self.default_model = self.module.SSDMobileNetV3Float32 + + def select_model(self, labels): + """Select best model for provided labels + Raises InvalidLabelException if any labels are unsupported + + Args: + labels: List of labels or None. If labels are not provided, default to SSDMobileNet with COCO labels + """ + + def _select(cls_list): + for cls_str in cls_list: + predictor_cls = getattr(self.module, cls_str) + if predictor_cls.validate_labels(labels): + return predictor_cls + else: + continue + raise InvalidLabelException + + if len(labels) is 0: + return self.default_model + else: + return _select(self.model_list) + + def label_map(self): + return {x: getattr(self.module, x).LABELS for x in self.model_list} diff --git a/rpi_deep_pantilt/detect/ssd_mobilenet_v3_coco.py b/rpi_deep_pantilt/detect/ssd_mobilenet_v3_coco.py deleted file mode 100644 index e7b3997..0000000 --- a/rpi_deep_pantilt/detect/ssd_mobilenet_v3_coco.py +++ /dev/null @@ -1,364 +0,0 @@ -# Python -import logging -import pathlib -import os -import sys - -# lib -import numpy as np -from PIL import Image -import tensorflow as tf - -from rpi_deep_pantilt import __path__ as rpi_deep_pantilt_path -from rpi_deep_pantilt.detect.util.label import create_category_index_from_labelmap -from rpi_deep_pantilt.detect.util.visualization import visualize_boxes_and_labels_on_image_array - -LABELS = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', - 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'] - - -class SSDMobileNet_V3_Coco_EdgeTPU_Quant(object): - - EDGETPU_SHARED_LIB = 'libedgetpu.so.1' - PATH_TO_LABELS = rpi_deep_pantilt_path[0] + '/data/mscoco_label_map.pbtxt' - - def __init__( - self, - base_url='https://github.com/leigh-johnson/rpi-deep-pantilt/releases/download/v1.0.0/', - model_name='ssdlite_mobilenet_edgetpu_coco_quant', - input_shape=(320, 320), - min_score_thresh=0.50, - tflite_model_file='model_postprocessed_quantized_128_uint8_edgetpu.tflite' - ): - - self.base_url = base_url - self.model_name = model_name - self.model_file = model_name + '.tar.gz' - self.model_url = base_url + self.model_file - self.tflite_model_file = tflite_model_file - - self.model_dir = tf.keras.utils.get_file( - fname=self.model_file, - origin=self.model_url, - untar=True, - cache_subdir='models' - ) - - self.min_score_thresh = min_score_thresh - - self.model_path = os.path.splitext( - os.path.splitext(self.model_dir)[0] - )[0] + f'/{self.tflite_model_file}' - - try: - from tflite_runtime import interpreter as coral_tflite_interpreter - except ImportError as e: - logging.error(e) - logging.error('Please install Edge TPU tflite_runtime:') - logging.error( - '$ pip install https://dl.google.com/coral/python/tflite_runtime-2.1.0.post1-cp37-cp37m-linux_armv7l.whl') - sys.exit(1) - - self.tflite_interpreter = coral_tflite_interpreter.Interpreter( - model_path=self.model_path, - experimental_delegates=[ - tf.lite.experimental.load_delegate(self.EDGETPU_SHARED_LIB) - ] - ) - - self.tflite_interpreter.allocate_tensors() - - self.input_details = self.tflite_interpreter.get_input_details() - self.output_details = self.tflite_interpreter.get_output_details() - - self.category_index = create_category_index_from_labelmap( - self.PATH_TO_LABELS, use_display_name=True) - - logging.info( - f'loaded labels from {self.PATH_TO_LABELS} \n {self.category_index}') - - logging.info(f'initialized model {model_name} \n') - logging.info( - f'model inputs: {self.input_details} \n {self.input_details}') - logging.info( - f'model outputs: {self.output_details} \n {self.output_details}') - - def label_to_category_index(self, labels): - # @todo :trashfire: - return tuple(map( - lambda x: x['id'], - filter( - lambda x: x['name'] in labels, self.category_index.values() - ) - )) - - def label_display_name_by_idx(self, idx): - return self.category_index[idx]['display_name'] - - def filter_tracked(self, prediction, label_idxs): - ''' - Zero predictions not in list of label_idxs - ''' - return { - 'detection_boxes': prediction.get('detection_boxes'), - 'detection_classes': prediction.get('detection_classes'), - 'detection_scores': - np.array( - [v if prediction.get('detection_classes')[i] in label_idxs - else 0.0 - for i, v in enumerate(prediction.get('detection_scores'))]) - } - - def create_overlay(self, image_np, output_dict): - - image_np = image_np.copy() - - # draw bounding boxes - visualize_boxes_and_labels_on_image_array( - image_np, - output_dict['detection_boxes'], - output_dict['detection_classes'], - output_dict['detection_scores'], - self.category_index, - use_normalized_coordinates=True, - line_thickness=4, - min_score_thresh=self.min_score_thresh, - max_boxes_to_draw=3 - ) - - img = Image.fromarray(image_np) - - return img.tobytes() - - def predict(self, image): - ''' - image - np.array (3 RGB channels) - - returns - { - 'detection_classes': int64, - 'num_detections': int64 - 'detection_masks': ... - } - ''' - - image = np.asarray(image) - # normalize 0 - 255 RGB to values between (-1, 1) - #image = (image / 128.0) - 1 - - # The input needs to be a tensor, convert it using `tf.convert_to_tensor`. - - input_tensor = tf.convert_to_tensor(image, dtype=tf.uint8) - - # The model expects a batch of images, so add an axis with `tf.newaxis`. - input_tensor = input_tensor[tf.newaxis, ...] - - # Run inference - self.tflite_interpreter.set_tensor( - self.input_details[0]['index'], input_tensor) - - self.tflite_interpreter.invoke() - - # TFLite_Detection_PostProcess custom op node has four outputs: - # detection_boxes: a float32 tensor of shape [1, num_boxes, 4] with box - # locations - # detection_classes: a float32 tensor of shape [1, num_boxes] - # with class indices - # detection_scores: a float32 tensor of shape [1, num_boxes] - # with class scores - # num_boxes: a float32 tensor of size 1 containing the number of detected - # boxes - - # Without the PostProcessing ops, the graph has two outputs: - # 'raw_outputs/box_encodings': a float32 tensor of shape [1, num_anchors, 4] - # containing the encoded box predictions. - # 'raw_outputs/class_predictions': a float32 tensor of shape - # [1, num_anchors, num_classes] containing the class scores for each anchor - # after applying score conversion. - - box_data = tf.convert_to_tensor(self.tflite_interpreter.get_tensor( - self.output_details[0]['index'])) - class_data = tf.convert_to_tensor(self.tflite_interpreter.get_tensor( - self.output_details[1]['index'])) - score_data = tf.convert_to_tensor(self.tflite_interpreter.get_tensor( - self.output_details[2]['index'])) - num_detections = tf.convert_to_tensor(self.tflite_interpreter.get_tensor( - self.output_details[3]['index'])) - - class_data = tf.squeeze( - class_data, axis=[0]).numpy().astype(np.int64) + 1 - box_data = tf.squeeze(box_data, axis=[0]).numpy() - score_data = tf.squeeze(score_data, axis=[0]).numpy() - - return { - 'detection_boxes': box_data, - 'detection_classes': class_data, - 'detection_scores': score_data, - 'num_detections': len(num_detections) - } - - -class SSDMobileNet_V3_Small_Coco_PostProcessed(object): - - PATH_TO_LABELS = rpi_deep_pantilt_path[0] + '/data/mscoco_label_map.pbtxt' - - def __init__( - self, - base_url='https://github.com/leigh-johnson/rpi-deep-pantilt/releases/download/v1.0.0/', - model_name='ssd_mobilenet_v3_small_coco_2019_08_14', - input_shape=(320, 320), - min_score_thresh=0.6 - - ): - - self.base_url = base_url - self.model_name = model_name - self.model_file = model_name + '.tar.gz' - self.model_url = base_url + self.model_file - self.min_score_thresh = min_score_thresh - - self.model_dir = tf.keras.utils.get_file( - fname=self.model_name, - origin=self.model_url, - untar=True, - cache_subdir='models' - ) - - self.model_path = os.path.splitext( - os.path.splitext(self.model_dir)[0] - )[0] + '/model_postprocessed.tflite' - - self.tflite_interpreter = tf.lite.Interpreter( - model_path=self.model_path, - ) - self.tflite_interpreter.allocate_tensors() - - self.input_details = self.tflite_interpreter.get_input_details() - self.output_details = self.tflite_interpreter.get_output_details() - - self.category_index = create_category_index_from_labelmap( - self.PATH_TO_LABELS, use_display_name=True) - - logging.info( - f'loaded labels from {self.PATH_TO_LABELS} \n {self.category_index}') - - logging.info(f'initialized model {model_name} \n') - logging.info( - f'model inputs: {self.input_details} \n {self.input_details}') - logging.info( - f'model outputs: {self.output_details} \n {self.output_details}') - - def label_to_category_index(self, labels): - # @todo :trashfire: - return tuple(map( - lambda x: x['id'], - filter( - lambda x: x['name'] in labels, self.category_index.values() - ) - )) - - def label_display_name_by_idx(self, idx): - return self.category_index[idx]['display_name'] - - def filter_tracked(self, prediction, label_idxs): - ''' - Zero predictions not in list of label_idxs - ''' - return { - 'detection_boxes': prediction.get('detection_boxes'), - 'detection_classes': prediction.get('detection_classes'), - 'detection_scores': - np.array( - [v if prediction.get('detection_classes')[i] in label_idxs - else 0.0 - for i, v in enumerate(prediction.get('detection_scores'))]) - } - - def create_overlay(self, image_np, output_dict): - - image_np = image_np.copy() - - # draw bounding boxes - visualize_boxes_and_labels_on_image_array( - image_np, - output_dict['detection_boxes'], - output_dict['detection_classes'], - output_dict['detection_scores'], - self.category_index, - use_normalized_coordinates=True, - line_thickness=4, - min_score_thresh=self.min_score_thresh, - max_boxes_to_draw=3 - ) - - img = Image.fromarray(image_np) - - return img.tobytes() - - def predict(self, image): - ''' - image - np.array (3 RGB channels) - - returns - { - 'detection_classes': int64, - 'num_detections': int64 - 'detection_masks': ... - } - ''' - - image = np.asarray(image) - # normalize 0 - 255 RGB to values between (-1, 1) - image = (image / 128.0) - 1 - - # The input needs to be a tensor, convert it using `tf.convert_to_tensor`. - - input_tensor = tf.convert_to_tensor(image, dtype=tf.float32) - - # The model expects a batch of images, so add an axis with `tf.newaxis`. - input_tensor = input_tensor[tf.newaxis, ...] - - # Run inference - self.tflite_interpreter.set_tensor( - self.input_details[0]['index'], input_tensor) - - self.tflite_interpreter.invoke() - - # TFLite_Detection_PostProcess custom op node has four outputs: - # detection_boxes: a float32 tensor of shape [1, num_boxes, 4] with box - # locations - # detection_classes: a float32 tensor of shape [1, num_boxes] - # with class indices - # detection_scores: a float32 tensor of shape [1, num_boxes] - # with class scores - # num_boxes: a float32 tensor of size 1 containing the number of detected - # boxes - - # Without the PostProcessing ops, the graph has two outputs: - # 'raw_outputs/box_encodings': a float32 tensor of shape [1, num_anchors, 4] - # containing the encoded box predictions. - # 'raw_outputs/class_predictions': a float32 tensor of shape - # [1, num_anchors, num_classes] containing the class scores for each anchor - # after applying score conversion. - - box_data = tf.convert_to_tensor(self.tflite_interpreter.get_tensor( - self.output_details[0]['index'])) - class_data = tf.convert_to_tensor(self.tflite_interpreter.get_tensor( - self.output_details[1]['index'])) - score_data = tf.convert_to_tensor(self.tflite_interpreter.get_tensor( - self.output_details[2]['index'])) - num_detections = tf.convert_to_tensor(self.tflite_interpreter.get_tensor( - self.output_details[3]['index'])) - - # hilarious, but it seems like all classes predictions are off by 1 idx - class_data = tf.squeeze( - class_data, axis=[0]).numpy().astype(np.int64) + 1 - box_data = tf.squeeze(box_data, axis=[0]).numpy() - score_data = tf.squeeze(score_data, axis=[0]).numpy() - - return { - 'detection_boxes': box_data, # 10, 4 - 'detection_classes': class_data, # 10 - 'detection_scores': score_data, # 10, - 'num_detections': len(num_detections) - } diff --git a/rpi_deep_pantilt/detect/util/exceptions.py b/rpi_deep_pantilt/detect/util/exceptions.py new file mode 100644 index 0000000..f30e29d --- /dev/null +++ b/rpi_deep_pantilt/detect/util/exceptions.py @@ -0,0 +1,2 @@ +class InvalidLabelException(Exception): + pass diff --git a/rpi_deep_pantilt/detect/util/label.py b/rpi_deep_pantilt/detect/util/label.py index 1012d84..604dee1 100644 --- a/rpi_deep_pantilt/detect/util/label.py +++ b/rpi_deep_pantilt/detect/util/label.py @@ -23,9 +23,7 @@ from rpi_deep_pantilt.detect.util import string_int_label_map_pb2 -def convert_label_map_to_categories(label_map, - max_num_classes, - use_display_name=True): +def convert_label_map_to_categories(label_map, max_num_classes, use_display_name=True): """Given label map proto returns categories list compatible with eval. This function converts label map proto and returns a list of dicts, each of @@ -54,24 +52,27 @@ def convert_label_map_to_categories(label_map, if not label_map: label_id_offset = 1 for class_id in range(max_num_classes): - categories.append({ - 'id': class_id + label_id_offset, - 'name': 'category_{}'.format(class_id + label_id_offset) - }) + categories.append( + { + "id": class_id + label_id_offset, + "name": "category_{}".format(class_id + label_id_offset), + } + ) return categories for item in label_map.item: if not 0 < item.id <= max_num_classes: logging.info( - 'Ignore item %d since it falls outside of requested ' - 'label range.', item.id) + "Ignore item %d since it falls outside of requested " "label range.", + item.id, + ) continue - if use_display_name and item.HasField('display_name'): + if use_display_name and item.HasField("display_name"): name = item.display_name else: name = item.name if item.id not in list_of_ids_already_added: list_of_ids_already_added.append(item.id) - categories.append({'id': item.id, 'name': name}) + categories.append({"id": item.id, "name": name}) return categories @@ -86,11 +87,13 @@ def _validate_label_map(label_map): """ for item in label_map.item: if item.id < 0: - raise ValueError('Label map ids should be >= 0.') - if (item.id == 0 and item.name != 'background' and - item.display_name != 'background'): - raise ValueError( - 'Label map id 0 is reserved for the background label') + raise ValueError("Label map ids should be >= 0.") + if ( + item.id == 0 + and item.name != "background" + and item.display_name != "background" + ): + raise ValueError("Label map id 0 is reserved for the background label") def load_labelmap(path): @@ -101,7 +104,7 @@ def load_labelmap(path): Returns: a StringIntLabelMapProto """ - with tf.compat.v1.gfile.GFile(path, 'r') as fid: + with tf.compat.v1.gfile.GFile(path, "r") as fid: label_map_string = fid.read() label_map = string_int_label_map_pb2.StringIntLabelMap() try: @@ -131,8 +134,7 @@ def create_categories_from_labelmap(label_map_path, use_display_name=True): """ label_map = load_labelmap(label_map_path) max_num_classes = max(item.id for item in label_map.item) - return convert_label_map_to_categories(label_map, max_num_classes, - use_display_name) + return convert_label_map_to_categories(label_map, max_num_classes, use_display_name) def create_category_index(categories): @@ -150,7 +152,7 @@ def create_category_index(categories): """ category_index = {} for cat in categories: - category_index[cat['id']] = cat + category_index[cat["id"]] = cat return category_index @@ -168,6 +170,5 @@ def create_category_index_from_labelmap(label_map_path, use_display_name=True): containing categories, e.g. {1: {'id': 1, 'name': 'dog'}, 2: {'id': 2, 'name': 'cat'}, ...} """ - categories = create_categories_from_labelmap( - label_map_path, use_display_name) + categories = create_categories_from_labelmap(label_map_path, use_display_name) return create_category_index(categories) diff --git a/rpi_deep_pantilt/detect/util/string_int_label_map_pb2.py b/rpi_deep_pantilt/detect/util/string_int_label_map_pb2.py index b9a7e42..3816b41 100644 --- a/rpi_deep_pantilt/detect/util/string_int_label_map_pb2.py +++ b/rpi_deep_pantilt/detect/util/string_int_label_map_pb2.py @@ -7,114 +7,162 @@ from google.protobuf import message as _message from google.protobuf import descriptor as _descriptor import sys -_b = sys.version_info[0] < 3 and ( - lambda x: x) or (lambda x: x.encode('latin1')) + +_b = sys.version_info[0] < 3 and (lambda x: x) or (lambda x: x.encode("latin1")) # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() DESCRIPTOR = _descriptor.FileDescriptor( - name='object_detection/protos/string_int_label_map.proto', - package='object_detection.protos', - syntax='proto2', + name="object_detection/protos/string_int_label_map.proto", + package="object_detection.protos", + syntax="proto2", serialized_options=None, - serialized_pb=_b('\n2object_detection/protos/string_int_label_map.proto\x12\x17object_detection.protos\"G\n\x15StringIntLabelMapItem\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\n\n\x02id\x18\x02 \x01(\x05\x12\x14\n\x0c\x64isplay_name\x18\x03 \x01(\t\"Q\n\x11StringIntLabelMap\x12<\n\x04item\x18\x01 \x03(\x0b\x32..object_detection.protos.StringIntLabelMapItem') + serialized_pb=_b( + '\n2object_detection/protos/string_int_label_map.proto\x12\x17object_detection.protos"G\n\x15StringIntLabelMapItem\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\n\n\x02id\x18\x02 \x01(\x05\x12\x14\n\x0c\x64isplay_name\x18\x03 \x01(\t"Q\n\x11StringIntLabelMap\x12<\n\x04item\x18\x01 \x03(\x0b\x32..object_detection.protos.StringIntLabelMapItem' + ), ) _STRINGINTLABELMAPITEM = _descriptor.Descriptor( - name='StringIntLabelMapItem', - full_name='object_detection.protos.StringIntLabelMapItem', + name="StringIntLabelMapItem", + full_name="object_detection.protos.StringIntLabelMapItem", filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( - name='name', full_name='object_detection.protos.StringIntLabelMapItem.name', index=0, - number=1, type=9, cpp_type=9, label=1, - has_default_value=False, default_value=_b("").decode('utf-8'), - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), + name="name", + full_name="object_detection.protos.StringIntLabelMapItem.name", + index=0, + number=1, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), _descriptor.FieldDescriptor( - name='id', full_name='object_detection.protos.StringIntLabelMapItem.id', index=1, - number=2, type=5, cpp_type=1, label=1, - has_default_value=False, default_value=0, - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), + name="id", + full_name="object_detection.protos.StringIntLabelMapItem.id", + index=1, + number=2, + type=5, + cpp_type=1, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), _descriptor.FieldDescriptor( - name='display_name', full_name='object_detection.protos.StringIntLabelMapItem.display_name', index=2, - number=3, type=9, cpp_type=9, label=1, - has_default_value=False, default_value=_b("").decode('utf-8'), - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - ], - extensions=[ + name="display_name", + full_name="object_detection.protos.StringIntLabelMapItem.display_name", + index=2, + number=3, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), ], + extensions=[], nested_types=[], - enum_types=[ - ], + enum_types=[], serialized_options=None, is_extendable=False, - syntax='proto2', + syntax="proto2", extension_ranges=[], - oneofs=[ - ], + oneofs=[], serialized_start=79, serialized_end=150, ) _STRINGINTLABELMAP = _descriptor.Descriptor( - name='StringIntLabelMap', - full_name='object_detection.protos.StringIntLabelMap', + name="StringIntLabelMap", + full_name="object_detection.protos.StringIntLabelMap", filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( - name='item', full_name='object_detection.protos.StringIntLabelMap.item', index=0, - number=1, type=11, cpp_type=10, label=3, - has_default_value=False, default_value=[], - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - ], - extensions=[ + name="item", + full_name="object_detection.protos.StringIntLabelMap.item", + index=0, + number=1, + type=11, + cpp_type=10, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), ], + extensions=[], nested_types=[], - enum_types=[ - ], + enum_types=[], serialized_options=None, is_extendable=False, - syntax='proto2', + syntax="proto2", extension_ranges=[], - oneofs=[ - ], + oneofs=[], serialized_start=152, serialized_end=233, ) -_STRINGINTLABELMAP.fields_by_name['item'].message_type = _STRINGINTLABELMAPITEM -DESCRIPTOR.message_types_by_name['StringIntLabelMapItem'] = _STRINGINTLABELMAPITEM -DESCRIPTOR.message_types_by_name['StringIntLabelMap'] = _STRINGINTLABELMAP +_STRINGINTLABELMAP.fields_by_name["item"].message_type = _STRINGINTLABELMAPITEM +DESCRIPTOR.message_types_by_name["StringIntLabelMapItem"] = _STRINGINTLABELMAPITEM +DESCRIPTOR.message_types_by_name["StringIntLabelMap"] = _STRINGINTLABELMAP _sym_db.RegisterFileDescriptor(DESCRIPTOR) -StringIntLabelMapItem = _reflection.GeneratedProtocolMessageType('StringIntLabelMapItem', (_message.Message,), { - 'DESCRIPTOR': _STRINGINTLABELMAPITEM, - '__module__': 'object_detection.protos.string_int_label_map_pb2' - # @@protoc_insertion_point(class_scope:object_detection.protos.StringIntLabelMapItem) -}) +StringIntLabelMapItem = _reflection.GeneratedProtocolMessageType( + "StringIntLabelMapItem", + (_message.Message,), + { + "DESCRIPTOR": _STRINGINTLABELMAPITEM, + "__module__": "object_detection.protos.string_int_label_map_pb2" + # @@protoc_insertion_point(class_scope:object_detection.protos.StringIntLabelMapItem) + }, +) _sym_db.RegisterMessage(StringIntLabelMapItem) -StringIntLabelMap = _reflection.GeneratedProtocolMessageType('StringIntLabelMap', (_message.Message,), { - 'DESCRIPTOR': _STRINGINTLABELMAP, - '__module__': 'object_detection.protos.string_int_label_map_pb2' - # @@protoc_insertion_point(class_scope:object_detection.protos.StringIntLabelMap) -}) +StringIntLabelMap = _reflection.GeneratedProtocolMessageType( + "StringIntLabelMap", + (_message.Message,), + { + "DESCRIPTOR": _STRINGINTLABELMAP, + "__module__": "object_detection.protos.string_int_label_map_pb2" + # @@protoc_insertion_point(class_scope:object_detection.protos.StringIntLabelMap) + }, +) _sym_db.RegisterMessage(StringIntLabelMap) diff --git a/rpi_deep_pantilt/detect/util/visualization.py b/rpi_deep_pantilt/detect/util/visualization.py index e20b08d..2ed7262 100644 --- a/rpi_deep_pantilt/detect/util/visualization.py +++ b/rpi_deep_pantilt/detect/util/visualization.py @@ -1,4 +1,3 @@ - # Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -26,29 +25,132 @@ import six STANDARD_COLORS = [ - 'AliceBlue', 'Chartreuse', 'Aqua', 'Aquamarine', 'Azure', 'Beige', 'Bisque', - 'BlanchedAlmond', 'BlueViolet', 'BurlyWood', 'CadetBlue', 'AntiqueWhite', - 'Chocolate', 'Coral', 'CornflowerBlue', 'Cornsilk', 'Crimson', 'Cyan', - 'DarkCyan', 'DarkGoldenRod', 'DarkGrey', 'DarkKhaki', 'DarkOrange', - 'DarkOrchid', 'DarkSalmon', 'DarkSeaGreen', 'DarkTurquoise', 'DarkViolet', - 'DeepPink', 'DeepSkyBlue', 'DodgerBlue', 'FireBrick', 'FloralWhite', - 'ForestGreen', 'Fuchsia', 'Gainsboro', 'GhostWhite', 'Gold', 'GoldenRod', - 'Salmon', 'Tan', 'HoneyDew', 'HotPink', 'IndianRed', 'Ivory', 'Khaki', - 'Lavender', 'LavenderBlush', 'LawnGreen', 'LemonChiffon', 'LightBlue', - 'LightCoral', 'LightCyan', 'LightGoldenRodYellow', 'LightGray', 'LightGrey', - 'LightGreen', 'LightPink', 'LightSalmon', 'LightSeaGreen', 'LightSkyBlue', - 'LightSlateGray', 'LightSlateGrey', 'LightSteelBlue', 'LightYellow', 'Lime', - 'LimeGreen', 'Linen', 'Magenta', 'MediumAquaMarine', 'MediumOrchid', - 'MediumPurple', 'MediumSeaGreen', 'MediumSlateBlue', 'MediumSpringGreen', - 'MediumTurquoise', 'MediumVioletRed', 'MintCream', 'MistyRose', 'Moccasin', - 'NavajoWhite', 'OldLace', 'Olive', 'OliveDrab', 'Orange', 'OrangeRed', - 'Orchid', 'PaleGoldenRod', 'PaleGreen', 'PaleTurquoise', 'PaleVioletRed', - 'PapayaWhip', 'PeachPuff', 'Peru', 'Pink', 'Plum', 'PowderBlue', 'Purple', - 'Red', 'RosyBrown', 'RoyalBlue', 'SaddleBrown', 'Green', 'SandyBrown', - 'SeaGreen', 'SeaShell', 'Sienna', 'Silver', 'SkyBlue', 'SlateBlue', - 'SlateGray', 'SlateGrey', 'Snow', 'SpringGreen', 'SteelBlue', 'GreenYellow', - 'Teal', 'Thistle', 'Tomato', 'Turquoise', 'Violet', 'Wheat', 'White', - 'WhiteSmoke', 'Yellow', 'YellowGreen' + "AliceBlue", + "Chartreuse", + "Aqua", + "Aquamarine", + "Azure", + "Beige", + "Bisque", + "BlanchedAlmond", + "BlueViolet", + "BurlyWood", + "CadetBlue", + "AntiqueWhite", + "Chocolate", + "Coral", + "CornflowerBlue", + "Cornsilk", + "Crimson", + "Cyan", + "DarkCyan", + "DarkGoldenRod", + "DarkGrey", + "DarkKhaki", + "DarkOrange", + "DarkOrchid", + "DarkSalmon", + "DarkSeaGreen", + "DarkTurquoise", + "DarkViolet", + "DeepPink", + "DeepSkyBlue", + "DodgerBlue", + "FireBrick", + "FloralWhite", + "ForestGreen", + "Fuchsia", + "Gainsboro", + "GhostWhite", + "Gold", + "GoldenRod", + "Salmon", + "Tan", + "HoneyDew", + "HotPink", + "IndianRed", + "Ivory", + "Khaki", + "Lavender", + "LavenderBlush", + "LawnGreen", + "LemonChiffon", + "LightBlue", + "LightCoral", + "LightCyan", + "LightGoldenRodYellow", + "LightGray", + "LightGrey", + "LightGreen", + "LightPink", + "LightSalmon", + "LightSeaGreen", + "LightSkyBlue", + "LightSlateGray", + "LightSlateGrey", + "LightSteelBlue", + "LightYellow", + "Lime", + "LimeGreen", + "Linen", + "Magenta", + "MediumAquaMarine", + "MediumOrchid", + "MediumPurple", + "MediumSeaGreen", + "MediumSlateBlue", + "MediumSpringGreen", + "MediumTurquoise", + "MediumVioletRed", + "MintCream", + "MistyRose", + "Moccasin", + "NavajoWhite", + "OldLace", + "Olive", + "OliveDrab", + "Orange", + "OrangeRed", + "Orchid", + "PaleGoldenRod", + "PaleGreen", + "PaleTurquoise", + "PaleVioletRed", + "PapayaWhip", + "PeachPuff", + "Peru", + "Pink", + "Plum", + "PowderBlue", + "Purple", + "Red", + "RosyBrown", + "RoyalBlue", + "SaddleBrown", + "Green", + "SandyBrown", + "SeaGreen", + "SeaShell", + "Sienna", + "Silver", + "SkyBlue", + "SlateBlue", + "SlateGray", + "SlateGrey", + "Snow", + "SpringGreen", + "SteelBlue", + "GreenYellow", + "Teal", + "Thistle", + "Tomato", + "Turquoise", + "Violet", + "Wheat", + "White", + "WhiteSmoke", + "Yellow", + "YellowGreen", ] @@ -74,13 +176,13 @@ def _get_multiplier_for_color_randomness(): return 1 # Return the closest prime number to num_colors / 10. - abs_distance = [np.abs(num_colors / 10. - p) for p in prime_candidates] + abs_distance = [np.abs(num_colors / 10.0 - p) for p in prime_candidates] num_candidates = len(abs_distance) inds = [i for _, i in sorted(zip(abs_distance, range(num_candidates)))] return prime_candidates[inds[0]] -def draw_mask_on_image_array(image, mask, color='red', alpha=0.4): +def draw_mask_on_image_array(image, mask, color="red", alpha=0.4): """Draws mask on an image. Args: @@ -94,34 +196,39 @@ def draw_mask_on_image_array(image, mask, color='red', alpha=0.4): ValueError: On incorrect data type for image or masks. """ if image.dtype != np.uint8: - raise ValueError('`image` not of type np.uint8') + raise ValueError("`image` not of type np.uint8") if mask.dtype != np.uint8: - raise ValueError('`mask` not of type np.uint8') + raise ValueError("`mask` not of type np.uint8") if np.any(np.logical_and(mask != 1, mask != 0)): - raise ValueError('`mask` elements should be in [0, 1]') + raise ValueError("`mask` elements should be in [0, 1]") if image.shape[:2] != mask.shape: - raise ValueError('The image has spatial dimensions %s but the mask has ' - 'dimensions %s' % (image.shape[:2], mask.shape)) + raise ValueError( + "The image has spatial dimensions %s but the mask has " + "dimensions %s" % (image.shape[:2], mask.shape) + ) rgb = ImageColor.getrgb(color) pil_image = Image.fromarray(image) - solid_color = np.expand_dims( - np.ones_like(mask), axis=2) * np.reshape(list(rgb), [1, 1, 3]) - pil_solid_color = Image.fromarray(np.uint8(solid_color)).convert('RGBA') - pil_mask = Image.fromarray(np.uint8(255.0*alpha*mask)).convert('L') + solid_color = np.expand_dims(np.ones_like(mask), axis=2) * np.reshape( + list(rgb), [1, 1, 3] + ) + pil_solid_color = Image.fromarray(np.uint8(solid_color)).convert("RGBA") + pil_mask = Image.fromarray(np.uint8(255.0 * alpha * mask)).convert("L") pil_image = Image.composite(pil_solid_color, pil_image, pil_mask) - np.copyto(image, np.array(pil_image.convert('RGB'))) - - -def draw_bounding_box_on_image(image, - ymin, - xmin, - ymax, - xmax, - color='red', - thickness=4, - display_str_list=(), - use_normalized_coordinates=True): + np.copyto(image, np.array(pil_image.convert("RGB"))) + + +def draw_bounding_box_on_image( + image, + ymin, + xmin, + ymax, + xmax, + color="red", + thickness=4, + display_str_list=(), + use_normalized_coordinates=True, +): """Adds a bounding box to an image. Bounding box coordinates can be specified in either absolute (pixel) or @@ -149,14 +256,21 @@ def draw_bounding_box_on_image(image, draw = ImageDraw.Draw(image) im_width, im_height = image.size if use_normalized_coordinates: - (left, right, top, bottom) = (xmin * im_width, xmax * im_width, - ymin * im_height, ymax * im_height) + (left, right, top, bottom) = ( + xmin * im_width, + xmax * im_width, + ymin * im_height, + ymax * im_height, + ) else: (left, right, top, bottom) = (xmin, xmax, ymin, ymax) - draw.line([(left, top), (left, bottom), (right, bottom), - (right, top), (left, top)], width=thickness, fill=color) + draw.line( + [(left, top), (left, bottom), (right, bottom), (right, top), (left, top)], + width=thickness, + fill=color, + ) try: - font = ImageFont.truetype('arial.ttf', 24) + font = ImageFont.truetype("arial.ttf", 24) except IOError: font = ImageFont.load_default() @@ -176,26 +290,32 @@ def draw_bounding_box_on_image(image, text_width, text_height = font.getsize(display_str) margin = np.ceil(0.05 * text_height) draw.rectangle( - [(left, text_bottom - text_height - 2 * margin), (left + text_width, - text_bottom)], - fill=color) + [ + (left, text_bottom - text_height - 2 * margin), + (left + text_width, text_bottom), + ], + fill=color, + ) draw.text( (left + margin, text_bottom - text_height - margin), display_str, - fill='black', - font=font) + fill="black", + font=font, + ) text_bottom -= text_height - 2 * margin -def draw_bounding_box_on_image_array(image, - ymin, - xmin, - ymax, - xmax, - color='red', - thickness=4, - display_str_list=(), - use_normalized_coordinates=True): +def draw_bounding_box_on_image_array( + image, + ymin, + xmin, + ymax, + xmax, + color="red", + thickness=4, + display_str_list=(), + use_normalized_coordinates=True, +): """Adds a bounding box to an image (numpy array). Bounding box coordinates can be specified in either absolute (pixel) or @@ -215,18 +335,24 @@ def draw_bounding_box_on_image_array(image, ymin, xmin, ymax, xmax as relative to the image. Otherwise treat coordinates as absolute. """ - image_pil = Image.fromarray(np.uint8(image)).convert('RGB') - draw_bounding_box_on_image(image_pil, ymin, xmin, ymax, xmax, color, - thickness, display_str_list, - use_normalized_coordinates) + image_pil = Image.fromarray(np.uint8(image)).convert("RGB") + draw_bounding_box_on_image( + image_pil, + ymin, + xmin, + ymax, + xmax, + color, + thickness, + display_str_list, + use_normalized_coordinates, + ) np.copyto(image, np.array(image_pil)) -def draw_keypoints_on_image(image, - keypoints, - color='red', - radius=2, - use_normalized_coordinates=True): +def draw_keypoints_on_image( + image, keypoints, color="red", radius=2, use_normalized_coordinates=True +): """Draws keypoints on an image. Args: @@ -245,16 +371,19 @@ def draw_keypoints_on_image(image, keypoints_x = tuple([im_width * x for x in keypoints_x]) keypoints_y = tuple([im_height * y for y in keypoints_y]) for keypoint_x, keypoint_y in zip(keypoints_x, keypoints_y): - draw.ellipse([(keypoint_x - radius, keypoint_y - radius), - (keypoint_x + radius, keypoint_y + radius)], - outline=color, fill=color) - - -def draw_keypoints_on_image_array(image, - keypoints, - color='red', - radius=2, - use_normalized_coordinates=True): + draw.ellipse( + [ + (keypoint_x - radius, keypoint_y - radius), + (keypoint_x + radius, keypoint_y + radius), + ], + outline=color, + fill=color, + ) + + +def draw_keypoints_on_image_array( + image, keypoints, color="red", radius=2, use_normalized_coordinates=True +): """Draws keypoints on an image (numpy array). Args: @@ -265,31 +394,33 @@ def draw_keypoints_on_image_array(image, use_normalized_coordinates: if True (default), treat keypoint values as relative to the image. Otherwise treat them as absolute. """ - image_pil = Image.fromarray(np.uint8(image)).convert('RGB') - draw_keypoints_on_image(image_pil, keypoints, color, radius, - use_normalized_coordinates) + image_pil = Image.fromarray(np.uint8(image)).convert("RGB") + draw_keypoints_on_image( + image_pil, keypoints, color, radius, use_normalized_coordinates + ) np.copyto(image, np.array(image_pil)) def visualize_boxes_and_labels_on_image_array( - image, - boxes, - classes, - scores, - category_index, - instance_masks=None, - instance_boundaries=None, - keypoints=None, - track_ids=None, - use_normalized_coordinates=False, - max_boxes_to_draw=20, - min_score_thresh=.5, - agnostic_mode=False, - line_thickness=4, - groundtruth_box_visualization_color='black', - skip_scores=False, - skip_labels=False, - skip_track_ids=False): + image, + boxes, + classes, + scores, + category_index, + instance_masks=None, + instance_boundaries=None, + keypoints=None, + track_ids=None, + use_normalized_coordinates=False, + max_boxes_to_draw=20, + min_score_thresh=0.5, + agnostic_mode=False, + line_thickness=4, + groundtruth_box_visualization_color="black", + skip_scores=False, + skip_labels=False, + skip_track_ids=False, +): """Overlay labeled boxes on an image with formatted scores and label names. This function groups boxes that correspond to the same location @@ -358,52 +489,47 @@ def visualize_boxes_and_labels_on_image_array( if scores is None: box_to_color_map[box] = groundtruth_box_visualization_color else: - display_str = '' + display_str = "" if not skip_labels: if not agnostic_mode: if classes[i] in six.viewkeys(category_index): - class_name = category_index[classes[i]]['name'] + class_name = category_index[classes[i]]["name"] else: - class_name = 'N/A' + class_name = "N/A" display_str = str(class_name) if not skip_scores: if not display_str: - display_str = '{}%'.format(int(100*scores[i])) + display_str = "{}%".format(int(100 * scores[i])) else: - display_str = '{}: {}%'.format( - display_str, int(100*scores[i])) + display_str = "{}: {}%".format( + display_str, int(100 * scores[i]) + ) if not skip_track_ids and track_ids is not None: if not display_str: - display_str = 'ID {}'.format(track_ids[i]) + display_str = "ID {}".format(track_ids[i]) else: - display_str = '{}: ID {}'.format( - display_str, track_ids[i]) + display_str = "{}: ID {}".format(display_str, track_ids[i]) box_to_display_str_map[box].append(display_str) if agnostic_mode: - box_to_color_map[box] = 'DarkOrange' + box_to_color_map[box] = "DarkOrange" elif track_ids is not None: prime_multipler = _get_multiplier_for_color_randomness() box_to_color_map[box] = STANDARD_COLORS[ - (prime_multipler * track_ids[i]) % len(STANDARD_COLORS)] + (prime_multipler * track_ids[i]) % len(STANDARD_COLORS) + ] else: box_to_color_map[box] = STANDARD_COLORS[ - classes[i] % len(STANDARD_COLORS)] + classes[i] % len(STANDARD_COLORS) + ] # Draw all boxes onto image. for box, color in box_to_color_map.items(): ymin, xmin, ymax, xmax = box if instance_masks is not None: - draw_mask_on_image_array( - image, - box_to_instance_masks_map[box], - color=color - ) + draw_mask_on_image_array(image, box_to_instance_masks_map[box], color=color) if instance_boundaries is not None: draw_mask_on_image_array( - image, - box_to_instance_boundaries_map[box], - color='red', - alpha=1.0 + image, box_to_instance_boundaries_map[box], color="red", alpha=1.0 ) draw_bounding_box_on_image_array( image, @@ -414,13 +540,15 @@ def visualize_boxes_and_labels_on_image_array( color=color, thickness=line_thickness, display_str_list=box_to_display_str_map[box], - use_normalized_coordinates=use_normalized_coordinates) + use_normalized_coordinates=use_normalized_coordinates, + ) if keypoints is not None: draw_keypoints_on_image_array( image, box_to_keypoints_map[box], color=color, radius=line_thickness / 2, - use_normalized_coordinates=use_normalized_coordinates) + use_normalized_coordinates=use_normalized_coordinates, + ) return image diff --git a/setup.cfg b/setup.cfg index a439b41..8867209 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,5 +1,5 @@ [bumpversion] -current_version = 1.2.1 +current_version = 2.0.0rc0 commit = True tag = True parse = (?P\d+)\.(?P\d+)\.(?P\d+) @@ -25,9 +25,6 @@ replace = __version__ = '{new_version}' [bdist_wheel] universal = 1 -[flake8] -exclude = docs - [aliases] test = pytest diff --git a/setup.py b/setup.py index f20c7b7..36b0070 100644 --- a/setup.py +++ b/setup.py @@ -111,7 +111,7 @@ def run(self): test_suite='tests', tests_require=test_requirements, url='https://github.com/leigh-johnson/rpi-deep-pantilt', - version='1.2.1', + version='2.0.0rc', zip_safe=False, ) diff --git a/tests/test_rpi_deep_pantilt.py b/tests/test_rpi_deep_pantilt.py index ff7a63d..3c1c3ab 100644 --- a/tests/test_rpi_deep_pantilt.py +++ b/tests/test_rpi_deep_pantilt.py @@ -32,7 +32,7 @@ def test_command_line_interface(): runner = CliRunner() result = runner.invoke(cli.main) assert result.exit_code == 0 - assert 'rpi_deep_pantilt.cli.main' in result.output - help_result = runner.invoke(cli.main, ['--help']) + assert "rpi_deep_pantilt.cli.main" in result.output + help_result = runner.invoke(cli.main, ["--help"]) assert help_result.exit_code == 0 - assert '--help Show this message and exit.' in help_result.output + assert "--help Show this message and exit." in help_result.output diff --git a/tools/install.sh b/tools/install.sh new file mode 100755 index 0000000..15d788c --- /dev/null +++ b/tools/install.sh @@ -0,0 +1,18 @@ +#!/bin/bash + +set +x + +echo "Installing platform dependencies" +sudo apt-get update && sudo apt-get install -y \ + cmake python3-dev libjpeg-dev libatlas-base-dev raspi-gpio libhdf5-dev python3-smbus libopenjp2-7-dev + +echo "Installing TensorFlow" +pip install https://github.com/leigh-johnson/Tensorflow-bin/releases/download/v2.2.0/tensorflow-2.2.0-cp37-cp37m-linux_armv7l.whl +echo "Installing Coral Edge TPU runtime" +pip install https://dl.google.com/coral/python/tflite_runtime-2.1.0.post1-cp37-cp37m-linux_armv7l.whl + +echo "Enabling Camera via raspi-config nonint" +sudo raspi-config nonint do_camera 1 + +echo "Enabling i2c dtparam via raspi-config nonint" +sudo raspi-config nonint do_i2c 1 \ No newline at end of file