mirror of
https://github.com/blakeblackshear/frigate.git
synced 2024-12-19 19:06:16 +01:00
Add features to rknn detector (#8631)
* 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
GitHub
parent
7b520e8a9d
commit
7522bb6fab
@ -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
|
||||
|
||||
Binary file not shown.
@ -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
|
||||
@ -326,3 +336,43 @@ model: # 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
|
||||
```
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user