Add label whitelist to filter detection types

DETECTION_LABELS env var accepts comma-separated list (e.g. "person,cat,dog").
Only matching detections are reported; others are ignored. Empty = report all.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Alex
2026-02-08 19:08:31 -06:00
parent b92c2064cd
commit 1bcf32889f
4 changed files with 17 additions and 0 deletions

View File

@@ -40,10 +40,12 @@ class ObjectDetector:
model_path: str,
labels_path: str,
confidence_threshold: float = 0.5,
label_whitelist: Optional[set[str]] = None,
):
self.model_path = Path(model_path)
self.labels_path = Path(labels_path)
self.confidence_threshold = confidence_threshold
self.label_whitelist = label_whitelist
self._interpreter = None
self._input_details = None
@@ -159,6 +161,10 @@ class ObjectDetector:
class_id = int(class_ids[i])
label = self._labels[class_id] if class_id < len(self._labels) else f"class_{class_id}"
# Skip labels not in whitelist (if set)
if self.label_whitelist and label not in self.label_whitelist:
continue
# Convert from (y_min, x_min, y_max, x_max) to (x_min, y_min, x_max, y_max)
y_min, x_min, y_max, x_max = boxes[i]
bbox = (