Mirrors/blakeblackshear.frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2024-12-19 19:06:16 +01:00
Code Issues Projects Releases Wiki Activity

Add features to rknn detector (#8631)

Browse Source 
* support for other yolov models and config checks

* apply code formatting

* Information about core mask and inference speed

* update rknn postprocess and remove params

* update model selection

* Apply suggestions from code review

Co-authored-by: Nicolas Mowen <nickmowen213@gmail.com>

---------

Co-authored-by: Nicolas Mowen <nickmowen213@gmail.com>
This commit is contained in:
Marc Altmann 2023-11-17 02:08:41 +01:00 committed by GitHub
parent 7b520e8a9d
commit 7522bb6fab
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 144 additions and 47 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
docker/rockchip/Dockerfile
View File

@ -9,10 +9,11 @@ COPY docker/rockchip/requirements-wheels-rk.txt /requirements-wheels-rk.txt
RUN sed -i "/https:\/\//d" /requirements-wheels.txt
RUN pip3 wheel --wheel-dir=/rk-wheels -c /requirements-wheels.txt -r /requirements-wheels-rk.txt
FROM wget as rk-libs
FROM wget as rk-downloads
RUN wget -qO librknnrt.so https://github.com/MarcA711/rknpu2/raw/master/runtime/RK3588/Linux/librknn_api/aarch64/librknnrt.so
RUN wget -qO ffmpeg https://github.com/MarcA711/Rockchip-FFmpeg-Builds/releases/download/latest/ffmpeg
RUN wget -qO ffprobe https://github.com/MarcA711/Rockchip-FFmpeg-Builds/releases/download/latest/ffprobe
RUN wget -qO yolov8n-320x320.rknn https://github.com/MarcA711/rknn-models/releases/download/latest/yolov8n-320x320.rknn
FROM deps AS rk-deps
ARG TARGETARCH
@ -21,12 +22,12 @@ RUN --mount=type=bind,from=rk-wheels,source=/rk-wheels,target=/deps/rk-wheels \
WORKDIR /opt/frigate/
COPY --from=rootfs / /
COPY --from=rk-libs /rootfs/librknnrt.so /usr/lib/
COPY docker/rockchip/yolov8n-320x320.rknn /models/
COPY --from=rk-downloads /rootfs/librknnrt.so /usr/lib/
COPY --from=rk-downloads /rootfs/yolov8n-320x320.rknn /models/
RUN rm -rf /usr/lib/btbn-ffmpeg/bin/ffmpeg
RUN rm -rf /usr/lib/btbn-ffmpeg/bin/ffprobe
COPY --from=rk-libs /rootfs/ffmpeg /usr/lib/btbn-ffmpeg/bin/
COPY --from=rk-libs /rootfs/ffprobe /usr/lib/btbn-ffmpeg/bin/
COPY --from=rk-downloads /rootfs/ffmpeg /usr/lib/btbn-ffmpeg/bin/
COPY --from=rk-downloads /rootfs/ffprobe /usr/lib/btbn-ffmpeg/bin/
RUN chmod +x /usr/lib/btbn-ffmpeg/bin/ffmpeg
RUN chmod +x /usr/lib/btbn-ffmpeg/bin/ffprobe

BIN
docker/rockchip/yolov8n-320x320.rknn
View File

Binary file not shown.

56
docs/docs/configuration/object_detectors.md
View File

@ -309,14 +309,24 @@ RKNN support is provided using the `-rk` suffix for the docker image. Moreover,
### Configuration
This `config.yml` shows all relevant options to configure the detector and explains them. All values shown are the default values (except for one). Lines that are required at least to use the detector are labeled as required, all other lines are optional.
```yaml
detectors: # required
rknn: # required
type: rknn # required
# core mask for npu
core_mask: 0
model: # required
# path to .rknn model file
path: /models/yolov8n-320x320.rknn
# name of yolov8 model or path to your own .rknn model file
# possible values are:
# - default-yolov8n
# - default-yolov8s
# - default-yolov8m
# - default-yolov8l
# - default-yolov8x
# - /config/model_cache/rknn/your_custom_model.rknn
path: default-yolov8n
# width and height of detection frames
width: 320
height: 320
@ -325,4 +335,44 @@ model: # required
input_pixel_format: bgr # required
# shape of detection frame
input_tensor: nhwc
```
```
Explanation for rknn specific options:
- **core mask** controls which cores of your NPU should be used. This option applies only to SoCs with a multicore NPU (at the time of writing this in only the RK3588/S). The easiest way is to pass the value as a binary number. To do so, use the prefix `0b` and write a `0` to disable a core and a `1` to enable a core, whereas the last digit coresponds to core0, the second last to core1, etc. You also have to use the cores in ascending order (so you can't use core0 and core2; but you can use core0 and core1). Enabling more cores can reduce the inference speed, especially when using bigger models (see section below). Examples:
- `core_mask: 0b000` or just `core_mask: 0` let the NPU decide which cores should be used. Default and recommended value.
- `core_mask: 0b001` use only core0.
- `core_mask: 0b011` use core0 and core1.
- `core_mask: 0b110` use core1 and core2. **This does not** work, since core0 is disabled.
### Choosing a model
There are 5 default yolov8 models that differ in size and therefore load the NPU more or less. In ascending order, with the top one being the smallest and least computationally intensive model:
| Model | Size in mb |
| ------- | ---------- |
| yolov8n | 9 |
| yolov8s | 25 |
| yolov8m | 54 |
| yolov8l | 90 |
| yolov8x | 136 |
:::tip
You can get the load of your NPU with the following command:
```bash
$ cat /sys/kernel/debug/rknpu/load
>> NPU load: Core0: 0%, Core1: 0%, Core2: 0%,
```
:::
- By default the rknn detector uses the yolov8n model (`model: path: default-yolov8n`). This model comes with the image, so no further steps than those mentioned above are necessary.
- If you want to use a more precise model, you can pass `default-yolov8s`, `default-yolov8m`, `default-yolov8l` or `default-yolov8x` as `model: path:` option.
- If the model does not exist, it will be automatically downloaded to `/config/model_cache/rknn`.
- If your server has no internet connection, you can download the model from [this Github repository](https://github.com/MarcA711/rknn-models/releases/tag/latest) using another device and place it in the `config/model_cache/rknn` on your system.
- Finally, you can also provide your own model. Note that only yolov8 models are currently supported. Moreover, you will need to convert your model to the rknn format using `rknn-toolkit2` on a x86 machine. Afterwards, you can place your `.rknn` model file in the `config/model_cache/rknn` directory on your system. Then you need to pass the path to your model using the `path` option of your `model` block like this:
```yaml
model:
path: /config/model_cache/rknn/my-rknn-model.rknn
```

124
frigate/detectors/plugins/rknn.py
View File

@ -1,8 +1,8 @@
import logging
import os.path
import urllib.request
from typing import Literal
import cv2
import cv2.dnn
import numpy as np
try:
@ -22,35 +22,83 @@ logger = logging.getLogger(__name__)
DETECTOR_KEY = "rknn"
yolov8_rknn_models = {
"default-yolov8n": "n",
"default-yolov8s": "s",
"default-yolov8m": "m",
"default-yolov8l": "l",
"default-yolov8x": "x",
}
class RknnDetectorConfig(BaseDetectorConfig):
type: Literal[DETECTOR_KEY]
score_thresh: float = Field(
default=0.5, ge=0, le=1, title="Minimal confidence for detection."
)
nms_thresh: float = Field(
default=0.45, ge=0, le=1, title="IoU threshold for non-maximum suppression."
)
core_mask: int = Field(default=0, ge=0, le=7, title="Core mask for NPU.")
class Rknn(DetectionApi):
type_key = DETECTOR_KEY
def __init__(self, config: RknnDetectorConfig):
self.model_path = config.model.path or "default-yolov8n"
self.core_mask = config.core_mask
self.height = config.model.height
self.width = config.model.width
self.score_thresh = config.score_thresh
self.nms_thresh = config.nms_thresh
self.model_path = config.model.path or "/models/yolov8n-320x320.rknn"
if self.model_path in yolov8_rknn_models:
if self.model_path == "default-yolov8n":
self.model_path = "/models/yolov8n-320x320.rknn"
else:
model_suffix = yolov8_rknn_models[self.model_path]
self.model_path = (
"/config/model_cache/rknn/yolov8{}-320x320.rknn".format(
model_suffix
)
)
os.makedirs("/config/model_cache/rknn", exist_ok=True)
if not os.path.isfile(self.model_path):
logger.info("Downloading yolov8{} model.".format(model_suffix))
urllib.request.urlretrieve(
"https://github.com/MarcA711/rknn-models/releases/download/latest/yolov8{}-320x320.rknn".format(
model_suffix
),
self.model_path,
)
if (config.model.width != 320) or (config.model.height != 320):
logger.error(
"Make sure to set the model width and heigth to 320 in your config.yml."
)
raise Exception(
"Make sure to set the model width and heigth to 320 in your config.yml."
)
if config.model.input_pixel_format != "bgr":
logger.error(
'Make sure to set the model input_pixel_format to "bgr" in your config.yml.'
)
raise Exception(
'Make sure to set the model input_pixel_format to "bgr" in your config.yml.'
)
if config.model.input_tensor != "nhwc":
logger.error(
'Make sure to set the model input_tensor to "nhwc" in your config.yml.'
)
raise Exception(
'Make sure to set the model input_tensor to "nhwc" in your config.yml.'
)
from rknnlite.api import RKNNLite
self.rknn = RKNNLite(verbose=False)
if self.rknn.load_rknn(self.model_path) != 0:
logger.error("Error initializing rknn model.")
if self.rknn.init_runtime() != 0:
logger.error("Error initializing rknn runtime.")
if self.rknn.init_runtime(core_mask=self.core_mask) != 0:
logger.error(
"Error initializing rknn runtime. Do you run docker in privileged mode?"
)
def __del__(self):
self.rknn.release()
@ -67,45 +115,43 @@ class Rknn(DetectionApi):
"""
results = np.transpose(results[0, :, :, 0]) # array shape (2100, 84)
classes = np.argmax(
results[:, 4:], axis=1
) # array shape (2100,); index of class with max confidence of each row
scores = np.max(
results[:, 4:], axis=1
) # array shape (2100,); max confidence of each row
# array shape (2100, 4); bounding box of each row
# remove lines with score scores < 0.4
filtered_arg = np.argwhere(scores > 0.4)
results = results[filtered_arg[:, 0]]
scores = scores[filtered_arg[:, 0]]
num_detections = len(scores)
if num_detections == 0:
return np.zeros((20, 6), np.float32)
if num_detections > 20:
top_arg = np.argpartition(scores, -20)[-20:]
results = results[top_arg]
scores = scores[top_arg]
num_detections = 20
classes = np.argmax(results[:, 4:], axis=1)
boxes = np.transpose(
np.vstack(
(
results[:, 0] - 0.5 * results[:, 2],
results[:, 1] - 0.5 * results[:, 3],
results[:, 2],
results[:, 3],
results[:, 0] - 0.5 * results[:, 2],
results[:, 3] + 0.5 * results[:, 3],
results[:, 2] + 0.5 * results[:, 2],
)
)
)
# indices of rows with confidence > SCORE_THRESH with Non-maximum Suppression (NMS)
result_boxes = cv2.dnn.NMSBoxes(
boxes, scores, self.score_thresh, self.nms_thresh, 0.5
)
detections = np.zeros((20, 6), np.float32)
for i in range(len(result_boxes)):
if i >= 20:
break
index = result_boxes[i]
detections[i] = [
classes[index],
scores[index],
(boxes[index][1]) / self.height,
(boxes[index][0]) / self.width,
(boxes[index][1] + boxes[index][3]) / self.height,
(boxes[index][0] + boxes[index][2]) / self.width,
]
detections[:num_detections, 0] = classes
detections[:num_detections, 1] = scores
detections[:num_detections, 2:] = boxes
return detections