modified: frigate/detectors/detector_config.py

modified:   frigate/detectors/plugins/cpu_tfl.py
	modified:   frigate/detectors/plugins/edgetpu_tfl.py
This commit is contained in:
Alec Groff 2024-01-29 09:16:04 -06:00
parent b1abdc2f8b
commit 61a808869b
3 changed files with 13 additions and 11 deletions

View File

@ -51,8 +51,8 @@ class ModelConfig(BaseModel):
model_type: ModelTypeEnum = Field( model_type: ModelTypeEnum = Field(
default=ModelTypeEnum.ssd, title="Object Detection Model Type" default=ModelTypeEnum.ssd, title="Object Detection Model Type"
) )
order: list[int] = Field( tfl_detector_output_tensor_order: list[int] = Field(
default=[0,1,2,3], title="Order Output Tensors [0=boxes,1=scores,2=class_ids,3=count]" default=[0,1,2,3], title="Order Output Tensors of TFL models [0=boxes,1=scores,2=class_ids,3=count]"
) )
_merged_labelmap: Optional[Dict[int, str]] = PrivateAttr() _merged_labelmap: Optional[Dict[int, str]] = PrivateAttr()
_colormap: Dict[int, Tuple[int, int, int]] = PrivateAttr() _colormap: Dict[int, Tuple[int, int, int]] = PrivateAttr()

View File

@ -37,15 +37,17 @@ class CpuTfl(DetectionApi):
self.tensor_input_details = self.interpreter.get_input_details() self.tensor_input_details = self.interpreter.get_input_details()
self.tensor_output_details = self.interpreter.get_output_details() self.tensor_output_details = self.interpreter.get_output_details()
self.tfl_detector_output_tensor_order = detector_config.model.tfl_detector_output_tensor_order
def detect_raw(self, tensor_input): def detect_raw(self, tensor_input):
self.interpreter.set_tensor(self.tensor_input_details[0]["index"], tensor_input) self.interpreter.set_tensor(self.tensor_input_details[0]["index"], tensor_input)
self.interpreter.invoke() self.interpreter.invoke()
boxes = self.interpreter.tensor(self.tensor_output_details[0]["index"])()[0] boxes = self.interpreter.tensor(self.tensor_output_details[self.tfl_detector_output_tensor_order[0]]["index"])()[0]
class_ids = self.interpreter.tensor(self.tensor_output_details[1]["index"])()[0] class_ids = self.interpreter.tensor(self.tensor_output_details[self.tfl_detector_output_tensor_order[1]]["index"])()[0]
scores = self.interpreter.tensor(self.tensor_output_details[2]["index"])()[0] scores = self.interpreter.tensor(self.tensor_output_details[self.tfl_detector_output_tensor_order[2]]["index"])()[0]
count = int( count = int(
self.interpreter.tensor(self.tensor_output_details[3]["index"])()[0] self.interpreter.tensor(self.tensor_output_details[self.tfl_detector_output_tensor_order[3]]["index"])()[0]
) )
detections = np.zeros((20, 6), np.float32) detections = np.zeros((20, 6), np.float32)

View File

@ -44,7 +44,7 @@ class EdgeTpuTfl(DetectionApi):
model_path=detector_config.model.path, model_path=detector_config.model.path,
experimental_delegates=[edge_tpu_delegate], experimental_delegates=[edge_tpu_delegate],
) )
self.order = detector_config.model.order self.tfl_detector_output_tensor_order = detector_config.model.tfl_detector_output_tensor_order
except ValueError: except ValueError:
logger.error( logger.error(
"No EdgeTPU was detected. If you do not have a Coral device yet, you must configure CPU detectors." "No EdgeTPU was detected. If you do not have a Coral device yet, you must configure CPU detectors."
@ -60,11 +60,11 @@ class EdgeTpuTfl(DetectionApi):
self.interpreter.set_tensor(self.tensor_input_details[0]["index"], tensor_input) self.interpreter.set_tensor(self.tensor_input_details[0]["index"], tensor_input)
self.interpreter.invoke() self.interpreter.invoke()
boxes = self.interpreter.tensor(self.tensor_output_details[self.order[0]]["index"])()[0] boxes = self.interpreter.tensor(self.tensor_output_details[self.tfl_detector_output_tensor_order[0]]["index"])()[0]
class_ids = self.interpreter.tensor(self.tensor_output_details[self.order[1]]["index"])()[0] class_ids = self.interpreter.tensor(self.tensor_output_details[self.tfl_detector_output_tensor_order[1]]["index"])()[0]
scores = self.interpreter.tensor(self.tensor_output_details[self.order[2]]["index"])()[0] scores = self.interpreter.tensor(self.tensor_output_details[self.tfl_detector_output_tensor_order[2]]["index"])()[0]
count = int( count = int(
self.interpreter.tensor(self.tensor_output_details[self.order[3]]["index"])()[0] self.interpreter.tensor(self.tensor_output_details[self.tfl_detector_output_tensor_order[3]]["index"])()[0]
) )
detections = np.zeros((20, 6), np.float32) detections = np.zeros((20, 6), np.float32)