remove 1-tensor processing. add pre_process() function

This commit is contained in:
Dan Brown 2025-12-02 16:01:58 +01:00
parent 1664b2f3bb
commit 6e288839be

View File

@ -75,14 +75,18 @@ class EdgeTpuTfl(DetectionApi):
self.min_score = 0.4
self.max_detections = 20
model_type = detector_config.model.model_type
self.model_type = detector_config.model.model_type
self.model_requires_int8 = self.tensor_input_details[0]["dtype"] == np.int8
if model_type == ModelTypeEnum.yologeneric
logger.debug(
f"Using YOLO postprocessing for {len(self.tensor_output_details)}-tensor output"
if self.model_type == ModelTypeEnum.yologeneric:
logger.debug("Using YOLO preprocessing/postprocessing")
if len(self.tensor_output_details) not in [2,3]:
logger.error(
f"Invalid count of output tensors in YOLO model. Found {len(self.tensor_output_details)}, expecting 2 or 3."
)
if len(self.tensor_output_details) > 1: # expecting 2 or 3
raise
self.reg_max = 16 # = 64 dfl_channels // 4 # YOLO standard
self.min_logit_value = np.log(
self.min_score / (1 - self.min_score)
@ -142,9 +146,9 @@ class EdgeTpuTfl(DetectionApi):
self.boxes_scale, self.boxes_zero_point = boxes_details["quantization"]
else:
if model_type not in [ModelTypeEnum.ssd, None]:
if self.model_type not in [ModelTypeEnum.ssd, None]:
logger.warning(
f"Unsupported model_type '{model_type}' for EdgeTPU detector, falling back to SSD"
f"Unsupported model_type '{self.model_type}' for EdgeTPU detector, falling back to SSD"
)
logger.debug("Using SSD preprocessing/postprocessing")
@ -202,36 +206,20 @@ class EdgeTpuTfl(DetectionApi):
else:
self.output_scores_index = index
def detect_raw(self, tensor_input):
def pre_process(self, tensor_input):
if self.model_requires_int8:
tensor_input = np.bitwise_xor(tensor_input, 128).view(
np.int8
) # shift by -128
return tensor_input
def detect_raw(self, tensor_input):
tensor_input = self.pre_process(tensor_input)
self.interpreter.set_tensor(self.tensor_input_details[0]["index"], tensor_input)
self.interpreter.invoke()
if model_type == ModelTypeEnum.yologeneric
output_tensor_count = len(self.tensor_output_details)
if output_tensor_count == 1:
# Single-tensor YOLO model
# model output is (1, NC+4, 2100) for 320x320 image size
# boxes as xywh (normalized to [0,1])
# followed by NC class probabilities (also [0,1])
# BEWARE the tensor has only one quantization scale/zero_point,
# so it should be assembled carefully to have a range of [0,1]
outputs = []
for output in self.tensor_output_details:
x = self.interpreter.get_tensor(output["index"])
scale, zero_point = output["quantization"]
x = (x.astype(np.float32) - zero_point) * scale
# Denormalize xywh by image size
x[:, [0, 2]] *= self.model_width
x[:, [1, 3]] *= self.model_height
outputs.append(x)
return post_process_yolo(outputs, self.model_width, self.model_height)
elif output_tensor_count in [2,3]:
if self.model_type == ModelTypeEnum.yologeneric:
# Multi-tensor YOLO model with (non-standard B(H*W)C output format).
# (the comments indicate the shape of tensors,
# using "2100" as the anchor count (for image size of 320x320),
@ -346,12 +334,6 @@ class EdgeTpuTfl(DetectionApi):
detections[:num_detections, 5] = final_boxes[:, 2] / self.model_width
return detections
else:
logger.error(
f"Invalid count of output tensors in YOLO model. Found {output_tensor_count}, expecting 1/2/3."
)
raise
else:
# Default SSD model
self.determine_indexes_for_non_yolo_models()