administrator/frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2026-02-07 11:45:24 +03:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

format rknn.py

Browse Source
This commit is contained in:
MarcA711 2023-11-01 19:28:10 +00:00
parent e0d9e7dceb
commit e3342cd2dd
1 changed files with 34 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -13,10 +13,16 @@ logger = logging.getLogger(__name__)
DETECTOR_KEY = "rknn" DETECTOR_KEY = "rknn"
class RknnDetectorConfig(BaseDetectorConfig): class RknnDetectorConfig(BaseDetectorConfig):
type: Literal[DETECTOR_KEY] type: Literal[DETECTOR_KEY]
score_thresh: float = Field(default=0.5, ge=0, le= 1, title="Minimal confidence for detection.") score_thresh: float = Field(
nms_thresh: float = Field(default=0.45, ge=0, le= 1, title="IoU threshold for non-maximum suppression.") 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."
)
class Rknn(DetectionApi): class Rknn(DetectionApi):
type_key = DETECTOR_KEY type_key = DETECTOR_KEY
@ -30,11 +36,12 @@ class Rknn(DetectionApi):
self.model_path = config.model.path or "/models/yolov8n-320x320.rknn" self.model_path = config.model.path or "/models/yolov8n-320x320.rknn"
from rknnlite.api import RKNNLite from rknnlite.api import RKNNLite
self.rknn = RKNNLite(verbose=False) self.rknn = RKNNLite(verbose=False)
if self.rknn.load_rknn(self.model_path) != 0: if self.rknn.load_rknn(self.model_path) != 0:
logger.error('Error initializing rknn model.') logger.error("Error initializing rknn model.")
if self.rknn.init_runtime() != 0: if self.rknn.init_runtime() != 0:
logger.error('Error initializing rknn runtime.') logger.error("Error initializing rknn runtime.")
def __del__(self): def __del__(self):
self.rknn.release() self.rknn.release()
@ -51,21 +58,29 @@ class Rknn(DetectionApi):
""" """
results = np.transpose(results[0, :, :, 0]) # array shape (2100, 84) 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 classes = np.argmax(
scores = np.max(results[:, 4:], axis=1) # array shape (2100,); max confidence of each row 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 # array shape (2100, 4); bounding box of each row
boxes = np.transpose( boxes = np.transpose(
np.vstack(( np.vstack(
(
results[:, 0] - 0.5 * results[:, 2], results[:, 0] - 0.5 * results[:, 2],
results[:, 1] - 0.5 * results[:, 3], results[:, 1] - 0.5 * results[:, 3],
results[:, 2], results[:, 2],
results[:,3]) results[:, 3],
)
) )
) )
# indices of rows with confidence > SCORE_THRESH with Non-maximum Suppression (NMS) # 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) result_boxes = cv2.dnn.NMSBoxes(
boxes, scores, self.score_thresh, self.nms_thresh, 0.5
)
detections = np.zeros((20, 6), np.float32) detections = np.zeros((20, 6), np.float32)
@ -80,7 +95,7 @@ class Rknn(DetectionApi):
(boxes[index][1]) / self.height, (boxes[index][1]) / self.height,
(boxes[index][0]) / self.width, (boxes[index][0]) / self.width,
(boxes[index][1] + boxes[index][3]) / self.height, (boxes[index][1] + boxes[index][3]) / self.height,
(boxes[index][0] + boxes[index][2]) / self.width (boxes[index][0] + boxes[index][2]) / self.width,
] ]
return detections return detections
@ -90,5 +105,9 @@ class Rknn(DetectionApi):
return self.rknn.inference(inputs=tensor_input) return self.rknn.inference(inputs=tensor_input)
def detect_raw(self, tensor_input): def detect_raw(self, tensor_input):
output = self.inference([tensor_input,]) output = self.inference(
[
tensor_input,
]
)
return self.postprocess(output[0]) return self.postprocess(output[0])