administrator/frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2026-02-07 03:35:26 +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

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

@ -13,10 +13,16 @@ logger = logging.getLogger(__name__)
DETECTOR_KEY = "rknn"
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.")
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."
)
class Rknn(DetectionApi):
type_key = DETECTOR_KEY
@ -30,11 +36,12 @@ class Rknn(DetectionApi):
self.model_path = config.model.path or "/models/yolov8n-320x320.rknn"
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.')
logger.error("Error initializing rknn model.")
if self.rknn.init_runtime() != 0:
logger.error('Error initializing rknn runtime.')
logger.error("Error initializing rknn runtime.")
def __del__(self):
self.rknn.release()
@ -50,22 +57,30 @@ class Rknn(DetectionApi):
detections: array with shape (20, 6) with 20 rows of (class, confidence, y_min, x_min, y_max, x_max)
"""
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
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
boxes = np.transpose(
np.vstack((
results[:,0] - 0.5 * results[:,2],
results[:,1] - 0.5 * results[:,3],
results[:,2],
results[:,3])
np.vstack(
(
results[:, 0] - 0.5 * results[:, 2],
results[:, 1] - 0.5 * results[:, 3],
results[:, 2],
results[:, 3],
)
)
)
# 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)
@ -80,7 +95,7 @@ class Rknn(DetectionApi):
(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
(boxes[index][0] + boxes[index][2]) / self.width,
]
return detections
@ -90,5 +105,9 @@ class Rknn(DetectionApi):
return self.rknn.inference(inputs=tensor_input)
def detect_raw(self, tensor_input):
output = self.inference([tensor_input,])
output = self.inference(
[
tensor_input,
]
)
return self.postprocess(output[0])