Post-process correct label push

This commit is contained in:
OmriAx 2025-03-03 12:57:04 +02:00
parent b1f2e299ee
commit 2bdbad72aa

View File

@ -324,30 +324,30 @@ class HailoDetector(DetectionApi):
threshold = 0.4 threshold = 0.4
all_detections = [] all_detections = []
# Loop over the output list (each element corresponds to one output stream) # Use the outer loop index to determine the class
for idx, detection_set in enumerate(infer_results): for class_id, detection_set in enumerate(infer_results):
# Skip empty arrays
if not isinstance(detection_set, np.ndarray) or detection_set.size == 0: if not isinstance(detection_set, np.ndarray) or detection_set.size == 0:
continue continue
logging.debug(f"[DETECT_RAW] Processing detection set {idx} with shape {detection_set.shape}") logging.debug(f"[DETECT_RAW] Processing detection set {class_id} with shape {detection_set.shape}")
# For each detection row in the set:
for det in detection_set: for det in detection_set:
# Expecting at least 5 elements: [ymin, xmin, ymax, xmax, confidence] # Expect at least 5 elements: [ymin, xmin, ymax, xmax, confidence]
if det.shape[0] < 5: if det.shape[0] < 5:
continue continue
score = float(det[4]) score = float(det[4])
if score < threshold: if score < threshold:
continue continue
# If there is a 6th element, assume it's a class id; otherwise use dummy class 0. # Instead of checking for a sixth element, use the outer index as the class
if det.shape[0] >= 6: cls = class_id
cls = int(det[5]) if hasattr(self, "labels") and self.labels:
logging.debug(f"[DETECT_RAW] Detected class id: {cls} -> {self.labels[cls]}")
print(f"[DETECT_RAW] Detected class id: {cls} -> {self.labels[cls]}")
else: else:
cls = 0 logging.debug(f"[DETECT_RAW] Detected class id: {cls}")
# Append in the order Frigate expects: [class_id, confidence, ymin, xmin, ymax, xmax] print(f"[DETECT_RAW] Detected class id: {cls}")
# Append in the order: [class_id, confidence, ymin, xmin, ymax, xmax]
all_detections.append([cls, score, det[0], det[1], det[2], det[3]]) all_detections.append([cls, score, det[0], det[1], det[2], det[3]])
# If no valid detections were found, return a zero array.
if len(all_detections) == 0: if len(all_detections) == 0:
return np.zeros((20, 6), dtype=np.float32) return np.zeros((20, 6), dtype=np.float32)