From 8ac3114f9a014356272b8a44b752e82df342f245 Mon Sep 17 00:00:00 2001
From: Nate Meyer <Nate.Devel@gmail.com>
Date: Fri, 6 Jan 2023 08:03:16 -0500
Subject: [PATCH] Cleanup Detector labelmap (#4932)

* Add missing labels to default labelmap.  Fill any holes with "unknown".  Remove unique labelmap for tensorrt.

* Replace "truck" with "car" on Openvino labelmap
---
 Dockerfile                           |  3 ++-
 docker/tensorrt_models.sh            |  3 ---
 docs/docs/configuration/detectors.md |  1 -
 frigate/util.py                      |  6 ++++--
 labelmap.txt                         | 13 ++++++++++++-
 5 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index eeabd6e44..6a804491a 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -91,7 +91,8 @@ RUN wget -qO cpu_model.tflite https://github.com/google-coral/test_data/raw/rele
 COPY labelmap.txt .
 # Copy OpenVino model
 COPY --from=ov-converter /models/public/ssdlite_mobilenet_v2/FP16 openvino-model
-RUN wget -q https://github.com/openvinotoolkit/open_model_zoo/raw/master/data/dataset_classes/coco_91cl_bkgr.txt -O openvino-model/coco_91cl_bkgr.txt
+RUN wget -q https://github.com/openvinotoolkit/open_model_zoo/raw/master/data/dataset_classes/coco_91cl_bkgr.txt -O openvino-model/coco_91cl_bkgr.txt && \
+    sed -i 's/truck/car/g' openvino-model/coco_91cl_bkgr.txt
 
 
 
diff --git a/docker/tensorrt_models.sh b/docker/tensorrt_models.sh
index fb685bf09..957e817d6 100755
--- a/docker/tensorrt_models.sh
+++ b/docker/tensorrt_models.sh
@@ -32,6 +32,3 @@ do
     python3 onnx_to_tensorrt.py -m ${model}
     cp /tensorrt_demos/yolo/${model}.trt ${OUTPUT_FOLDER}/${model}.trt;
 done
-
-# Download Labelmap
-wget -q https://github.com/openvinotoolkit/open_model_zoo/raw/master/data/dataset_classes/coco_91cl.txt -O ${OUTPUT_FOLDER}/coco_91cl.txt
\ No newline at end of file
diff --git a/docs/docs/configuration/detectors.md b/docs/docs/configuration/detectors.md
index b2478d690..c8e1cc508 100644
--- a/docs/docs/configuration/detectors.md
+++ b/docs/docs/configuration/detectors.md
@@ -228,7 +228,6 @@ detectors:
 
 model:
   path: /trt-models/yolov7-tiny-416.trt
-  labelmap_path: /trt-models/coco_91cl.txt
   input_tensor: nchw
   input_pixel_format: rgb
   width: 416
diff --git a/frigate/util.py b/frigate/util.py
index 79fe7932f..d8cdef0d9 100755
--- a/frigate/util.py
+++ b/frigate/util.py
@@ -706,15 +706,17 @@ def load_labels(path, encoding="utf-8"):
       Dictionary mapping indices to labels.
     """
     with open(path, "r", encoding=encoding) as f:
+        labels = {index: "unknown" for index in range(91)}
         lines = f.readlines()
         if not lines:
             return {}
 
         if lines[0].split(" ", maxsplit=1)[0].isdigit():
             pairs = [line.split(" ", maxsplit=1) for line in lines]
-            return {int(index): label.strip() for index, label in pairs}
+            labels.update({int(index): label.strip() for index, label in pairs})
         else:
-            return {index: line.strip() for index, line in enumerate(lines)}
+            labels.update({index: line.strip() for index, line in enumerate(lines)})
+        return labels
 
 
 def clean_camera_user_pass(line: str) -> str:
diff --git a/labelmap.txt b/labelmap.txt
index a5a6fcb09..79fff1772 100644
--- a/labelmap.txt
+++ b/labelmap.txt
@@ -9,6 +9,7 @@
 8  boat
 9  traffic light
 10  fire hydrant
+11  street sign
 12  stop sign
 13  parking meter
 14  bench
@@ -22,8 +23,11 @@
 22  bear
 23  zebra
 24  giraffe
+25  hat
 26  backpack
 27  umbrella
+28  shoe
+29  eye glasses
 30  handbag
 31  tie
 32  suitcase
@@ -38,6 +42,7 @@
 41  surfboard
 42  tennis racket
 43  bottle
+44  plate
 45  wine glass
 46  cup
 47  fork
@@ -58,8 +63,12 @@
 62  couch
 63  potted plant
 64  bed
+65  mirror
 66  dining table
+67  window
+68  desk
 69  toilet
+70  door
 71  tv
 72  laptop
 73  mouse
@@ -71,10 +80,12 @@
 79  toaster
 80  sink
 81  refrigerator
+82  blender
 83  book
 84  clock
 85  vase
 86  scissors
 87  teddy bear
 88  hair drier
-89  toothbrush
\ No newline at end of file
+89  toothbrush
+90  hair brush
\ No newline at end of file