From 7a4eb0b37c1b68fe7bff5c1b05cb42d890ae5de8 Mon Sep 17 00:00:00 2001
From: Marc Altmann <40744649+MarcA711@users.noreply.github.com>
Date: Mon, 3 Jun 2024 03:47:26 +0200
Subject: [PATCH] Add coco-80 labelmap and update FFmpeg for Rockchip  (#11695)

* add coco-80 labelmap and update ffmpeg

* Update docs/docs/configuration/object_detectors.md

---------

Co-authored-by: Blake Blackshear <blake.blackshear@gmail.com>
---
 docker/main/rootfs/labelmap/coco-80.txt     | 80 +++++++++++++++++++++
 docker/rockchip/Dockerfile                  |  4 +-
 docs/docs/configuration/object_detectors.md |  8 +++
 frigate/detectors/plugins/rknn.py           | 10 +--
 4 files changed, 95 insertions(+), 7 deletions(-)
 create mode 100644 docker/main/rootfs/labelmap/coco-80.txt

diff --git a/docker/main/rootfs/labelmap/coco-80.txt b/docker/main/rootfs/labelmap/coco-80.txt
new file mode 100644
index 000000000..79e017175
--- /dev/null
+++ b/docker/main/rootfs/labelmap/coco-80.txt
@@ -0,0 +1,80 @@
+0  person
+1  bicycle
+2  car
+3  motorcycle
+4  airplane
+5  car
+6  train
+7  car
+8  boat
+9  traffic light
+10  fire hydrant
+11  stop sign
+12  parking meter
+13  bench
+14  bird
+15  cat
+16  dog
+17  horse
+18  sheep
+19  cow
+20  elephant
+21  bear
+22  zebra
+23  giraffe
+24  backpack
+25  umbrella
+26  handbag
+27  tie
+28  suitcase
+29  frisbee
+30  skis
+31  snowboard
+32  sports ball
+33  kite
+34  baseball bat
+35  baseball glove
+36  skateboard
+37  surfboard
+38  tennis racket
+39  bottle
+40  wine glass
+41  cup
+42  fork
+43  knife
+44  spoon
+45  bowl
+46  banana
+47  apple
+48  sandwich
+49  orange
+50  broccoli
+51  carrot
+52  hot dog
+53  pizza
+54  donut
+55  cake
+56  chair
+57  couch
+58  potted plant
+59  bed
+60  dining table
+61  toilet
+62  tv
+63  laptop
+64  mouse
+65  remote
+66  keyboard
+67  cell phone
+68  microwave
+69  oven
+70  toaster
+71  sink
+72  refrigerator
+73  book
+74  clock
+75  vase
+76  scissors
+77  teddy bear
+78  hair drier
+79  toothbrush
\ No newline at end of file
diff --git a/docker/rockchip/Dockerfile b/docker/rockchip/Dockerfile
index 7913857bd..e6222a67a 100644
--- a/docker/rockchip/Dockerfile
+++ b/docker/rockchip/Dockerfile
@@ -22,5 +22,5 @@ ADD https://github.com/MarcA711/rknn-toolkit2/releases/download/v2.0.0/librknnrt
 
 RUN rm -rf /usr/lib/btbn-ffmpeg/bin/ffmpeg
 RUN rm -rf /usr/lib/btbn-ffmpeg/bin/ffprobe
-ADD --chmod=111 https://github.com/MarcA711/Rockchip-FFmpeg-Builds/releases/download/6.1-3/ffmpeg /usr/lib/btbn-ffmpeg/bin/
-ADD --chmod=111 https://github.com/MarcA711/Rockchip-FFmpeg-Builds/releases/download/6.1-3/ffprobe /usr/lib/btbn-ffmpeg/bin/
+ADD --chmod=111 https://github.com/MarcA711/Rockchip-FFmpeg-Builds/releases/download/6.1-4/ffmpeg /usr/lib/btbn-ffmpeg/bin/
+ADD --chmod=111 https://github.com/MarcA711/Rockchip-FFmpeg-Builds/releases/download/6.1-4/ffprobe /usr/lib/btbn-ffmpeg/bin/
diff --git a/docs/docs/configuration/object_detectors.md b/docs/docs/configuration/object_detectors.md
index 9c4dc9893..f3a389798 100644
--- a/docs/docs/configuration/object_detectors.md
+++ b/docs/docs/configuration/object_detectors.md
@@ -348,8 +348,16 @@ model: # required
   input_pixel_format: bgr # required
   # shape of detection frame
   input_tensor: nhwc
+  # needs to be adjusted to model, see below
+  labelmap_path: /labelmap.txt # required
 ```
 
+The correct labelmap must be loaded for each model. If you use a custom model (see notes below), you must make sure to provide the correct labelmap. The table below lists the correct paths for the bundled models:
+
+| `path`              | `labelmap_path`       |
+| ------------------- | --------------------- |
+| deci-fp16-yolonas_* | /labelmap/coco-80.txt |
+
 ### Choosing a model
 
 :::warning
diff --git a/frigate/detectors/plugins/rknn.py b/frigate/detectors/plugins/rknn.py
index 03eca80b8..49f1f4c02 100644
--- a/frigate/detectors/plugins/rknn.py
+++ b/frigate/detectors/plugins/rknn.py
@@ -42,11 +42,11 @@ class Rknn(DetectionApi):
             config.model.model_type = model_props["model_type"]
 
         if model_props["model_type"] == ModelTypeEnum.yolonas:
-            logger.info("""
-                        You are using yolo-nas with weights from DeciAI.
-                        These weights are subject to their license and can't be used commercially.
-                        For more information, see: https://docs.deci.ai/super-gradients/latest/LICENSE.YOLONAS.html
-                        """)
+            logger.info(
+                "You are using yolo-nas with weights from DeciAI. "
+                "These weights are subject to their license and can't be used commercially. "
+                "For more information, see: https://docs.deci.ai/super-gradients/latest/LICENSE.YOLONAS.html"
+            )
 
         from rknnlite.api import RKNNLite