#patch: support tpu compatible docker image

# Please enter the commit message for your changes. Lines starting # with '#' will be kept; you may remove them yourself if you want to. # An empty message aborts the commit. # # Date: Fri Aug 4 16:51:36 2023 -0700 # # On branch main # Your branch and 'origin/main' have diverged, # and have 1 and 1 different commits each, respectively. # (use "git pull" to merge the remote branch into yours) # # Changes to be committed: # modified: .github/workflows/deploy.yaml # new file: Dockerfile.gpu # new file: requirements.gpu.txt # # Changes not staged for commit: # modified: .github/workflows/deploy.yaml # modified: Dockerfile # deleted: Dockerfile.gpu # deleted: requirements.gpu.txt # # Untracked files: # Dockerfile.tpu # requirements.tpu.txt #
circa10a · Aug 7, 2023 · aab5b75 · aab5b75
1 parent d9b42da
commit aab5b75
Show file tree

Hide file tree

Showing 8 changed files with 124 additions and 8 deletions.
diff --git a/.dockerignore b/.dockerignore
@@ -0,0 +1,2 @@
+__pycache__
+*.pt
diff --git a/.github/workflows/deploy.yaml b/.github/workflows/deploy.yaml
@@ -38,10 +38,43 @@ jobs:
         with:
           username: ${{ github.repository_owner }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
-      - name: Build and push
+      - name: Build and push CPU
         id: docker_build_push
         uses: docker/build-push-action@v4
         with:
           push: true
           tags: ${{ github.repository }}:latest,${{ github.repository }}:${{ needs.semver.outputs.tag }}
-          platforms: ${{ env.PLATFORMS }}
+          platforms: ${{ env.PLATFORMS }}
+  publish-docker-image-tpu:
+    runs-on: ubuntu-latest
+    needs: semver
+    steps:
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v2
+      - name: Set up Docker Build
+        uses: docker/setup-buildx-action@v2
+      - name: Login to DockerHub
+        uses: docker/login-action@v2
+        with:
+          username: ${{ github.repository_owner }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+      - name: Delete huge unnecessary tools folder
+        run: rm -rf /opt/hostedtoolcache
+      - name: Build and push arm64 TPU
+        uses: docker/build-push-action@v4
+        with:
+          push: true
+          file: ./Dockerfile.tpu
+          build-args: |
+            ULTRALYTICS_TAG=latest-arm64
+          tags: ${{ github.repository }}:latest-arm64-tpu,${{ github.repository }}:${{ needs.semver.outputs.tag }}-arm64-tpu
+          platforms: linux/arm64
+      - name: Build and push amd64 TPU
+        uses: docker/build-push-action@v4
+        with:
+          push: true
+          file: ./Dockerfile.tpu
+          build-args: |
+            ULTRALYTICS_TAG=latest-cpu
+          tags: ${{ github.repository }}:latest-amd64-tpu,${{ github.repository }}:${{ needs.semver.outputs.tag }}-amd64-tpu
+          platforms: linux/amd64
diff --git a/Dockerfile b/Dockerfile
@@ -4,7 +4,7 @@ WORKDIR $WORKDIR
 COPY . .
 ARG MODEL=yolov8s.pt
 RUN apt-get update && \
-    apt-get install -y curl libgl1-mesa-glx libglib2.0-0 libpython3-dev gnupg g++ && \
+    apt-get install -y libgl1-mesa-glx libglib2.0-0 libpython3-dev gnupg g++ libusb-1.0-0 && \
     apt-get upgrade -y && \
     apt-get autoremove -y && \
     rm -rf /var/lib/apt/lists/* && \

diff --git a/Dockerfile.tpu b/Dockerfile.tpu
@@ -0,0 +1,31 @@
+ARG ULTRALYTICS_TAG
+ARG MODEL=yolov8s.pt
+
+FROM ultralytics/ultralytics:latest-cpu AS EXPORT
+ARG MODEL
+ENV WORKDIR=/opt/cv-notifier
+WORKDIR $WORKDIR
+COPY . .
+RUN echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | tee /etc/apt/sources.list.d/coral-edgetpu.list && \
+    curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - && \
+    curl -sL https://github.com/ultralytics/assets/releases/download/v0.0.0/$MODEL -o $MODEL && \
+    apt-get update && \
+    DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y libedgetpu1-std edgetpu-compiler cmake libgl1-mesa-glx libglib2.0-0 libpython3-dev gnupg g++ libusb-1.0-0 tzdata && \
+    pip3 install -r requirements.tpu.build.txt --no-cache-dir && \
+    yolo export model=$MODEL format="edgetpu"
+
+FROM ultralytics/ultralytics:$ULTRALYTICS_TAG
+ENV WORKDIR=/opt/cv-notifier \
+    EDGETPU_ENABLED=true
+WORKDIR $WORKDIR
+COPY . .
+COPY --from=EXPORT $WORKDIR/yolov8s_saved_model/* .
+RUN echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | tee /etc/apt/sources.list.d/coral-edgetpu.list && \
+    curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - && \
+    apt-get update && \
+    DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y libedgetpu1-std libgl1-mesa-glx libglib2.0-0 libpython3-dev gnupg g++ libusb-1.0-0 tzdata && \
+    apt-get upgrade -y && \
+    apt-get autoremove -y && \
+    pip3 install -r requirements.tpu.txt --no-cache-dir
+
+CMD ["python3", "main.py"]
diff --git a/README.md b/README.md
@@ -12,6 +12,7 @@ Easily detect objects using computer vision ([YOLOv8](https://github.com/ultraly
   * [System requirements](#system-requirements)
   * [Local](#local)
   * [Docker](#docker)
+  * [Edge TPU](#edge-tpu)
 * [Configuration](#configuration)
   * [Variable substituiton](#variable-substituiton)
 
@@ -97,6 +98,41 @@ services:
       - /etc/localtime:/etc/localtime:ro
 ```
 
+### Edge TPU
+
+`cv-notifier` supports using edge TPU devices like the [Google Coral USB Accelerator](https://coral.ai/products/accelerator)
+
+Docker run:
+
+> Variable substitution in the configuration is optional. This is to demonstrate capabilities.
+
+```console
+docker run --name cv-notifier \
+    --privileged=true \
+    -v $PWD/sample.config.yaml:/opt/cv-notifier/config.yaml \
+    -v /etc/localtime:/etc/localtime:ro \
+    -v /dev/bus/usb:/dev/bus/usb
+    circa10a/cv-notifier:latest-tpu
+```
+
+Docker Compose:
+
+```yaml
+version: '3'
+
+services:
+  cv-notifier:
+    container_name: cv-notifier
+    privileged: true
+    restart: always
+    image: circa10a/cv-notifier:latest-tpu
+    env_file: .env
+    volumes:
+      - ./sample.config.yaml:/opt/cv-notifier/config.yaml
+      - /etc/localtime:/etc/localtime:ro
+      - /dev/bus/usb:/dev/bus/usb
+```
+
 ## Configuration
 
 `cv-notifier` supports many different configuration options. See [usage](#usage)
@@ -105,6 +141,7 @@ services:
 |-------------------------------------|----------------------------------------------------------------------|-----------|--------------------|--------------------------------|
 | Key                                 | Description                                                          | Required  | Default            | Supports environment variables |
 | `config.source`                     | Source of video stream such as `RTSP`, `RTMP`, `HTTP`                | `True`    | `None`             |  ✅                            |
+| `config.validate`                   | Validate objects in config are supported by model                    | `False`   | `True`             |  ❌                             |
 | `config.model`                      | What pre-trained model to use                                        | `False`   | `yolov8s.pt`       |  ❌                            |
 | `config.confidence`                 | Score filter to only show detected objects with confidence above this| `False`   | `0.50` (50%)       |  ❌                            |
 | `config.loglevel`                   | Level of logging. `info`, `debug`, `warning`                         | `False`   | `info`             |  ❌                            |
@@ -118,7 +155,6 @@ services:
 | `config.webhooks[0].headers`        | Map of HTTP headers to send in request                               | `False`   | `None`             |  ✅                            |
 | `config.webhooks[0].body`           | HTTP body to send in request                                         | `False`   | `None`             |  ✅                            |
 
-
 ### Variable substituiton
 
 Configuration fields that often require sensitive data do support environment variable subsition in the form of `$variable` within the configuration values. The `body` field within a webhook configuration supports some additional data around the object detected. This information can be used to customize request payloads like the example mentioned above in [configuration](#configuration).

diff --git a/lib/config/config.py b/lib/config/config.py
@@ -5,7 +5,7 @@
 from typing import List
 import yaml
 from loguru import logger as log
-from os import environ
+from os import getenv, environ
 from schema import Optional, Or, Regex, Schema
 from string import Template
 from ultralytics import YOLO
@@ -15,15 +15,18 @@
 
 # We'll use this map to store all objects to be recognized per webhook
 # then store the index of the webhooks that need to be called when a configured object is detected.
-# This keeps us from having to loop all webhooks/objects every time an object is detecteced.
+# This keeps us from having to loop all webhooks/objects every time an object is detected.
 OBJECT_TO_WEBHOOK_MAP = {}
 
+EDGETPU_ENABLED = getenv('EDGETPU_ENABLED', False).lower() in ('true', '1', 't')
+
 config_schema = Schema(
     {
         'config': {
             'source': str,
             Optional('confidence', default=0.5): float,
-            Optional('model', default='yolov8s.pt'): str,
+            Optional('validate', default=True): bool,
+            Optional('model', default=EDGETPU_ENABLED and 'yolov8s_full_integer_quant_edgetpu.tflite' or 'yolov8s.pt'): str,
             Optional('loglevel', default='info'): Or(
              'trace',
              'debug',
@@ -107,6 +110,7 @@ def parse(config_file: str) -> dict:
                 OBJECT_TO_WEBHOOK_MAP[object_name].append(webhook)
 
         # Validate objects in config are supported by model
-        validate_objects(model.names, OBJECT_TO_WEBHOOK_MAP.keys())
+        if config['config']['validate']:
+            validate_objects(model.names, OBJECT_TO_WEBHOOK_MAP.keys())
 
     return config
diff --git a/requirements.tpu.build.txt b/requirements.tpu.build.txt
@@ -0,0 +1,5 @@
+# Export tooling
+onnx >= 1.12.0 # ONNX export
+onnxsim >= 0.4.1 # ONNX simplifier
+tensorflow >= 2.4.1 # TF exports (-cpu, -aarch64, -macos)
+tflite-support >= 0.4.4
diff --git a/requirements.tpu.txt b/requirements.tpu.txt
@@ -0,0 +1,5 @@
+ultralytics >= 8.0.106
+schema >= 0.7.5
+requests >= 2.31.0
+loguru >= 0.7.0
+--extra-index-url https://google-coral.github.io/py-repo/ pycoral~=2.0