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:
@@ -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 = (
|
||||
|
||||
Reference in New Issue
Block a user