37 lines
1.3 KiB
Python
37 lines
1.3 KiB
Python
"""SAFE / PARTIAL / UNSAFE classification (helmet + vest focus)."""
|
|
from __future__ import annotations
|
|
|
|
from typing import Dict, List, Tuple
|
|
|
|
from saqr.core.detection import POSITIVE_TO_NEGATIVE, PPE_DISPLAY_ORDER
|
|
|
|
|
|
def status_from_items(items: Dict[str, float]) -> str:
|
|
has_helmet = items.get("helmet", 0.0) > items.get("no-helmet", 0.0) and items.get("helmet", 0.0) > 0
|
|
has_vest = items.get("vest", 0.0) > items.get("no-vest", 0.0) and items.get("vest", 0.0) > 0
|
|
no_helmet = items.get("no-helmet", 0.0) > 0
|
|
no_vest = items.get("no-vest", 0.0) > 0
|
|
|
|
if no_helmet or no_vest:
|
|
return "UNSAFE"
|
|
if has_helmet and has_vest:
|
|
return "SAFE"
|
|
if has_helmet or has_vest:
|
|
return "PARTIAL"
|
|
return "UNSAFE"
|
|
|
|
|
|
def split_wearing_missing(items: Dict[str, float]) -> Tuple[List[str], List[str], List[str]]:
|
|
wearing, missing, unknown = [], [], []
|
|
for pos in PPE_DISPLAY_ORDER:
|
|
neg = POSITIVE_TO_NEGATIVE[pos]
|
|
pos_conf = items.get(pos, 0.0)
|
|
neg_conf = items.get(neg, 0.0)
|
|
if pos_conf > neg_conf and pos_conf > 0:
|
|
wearing.append(pos)
|
|
elif neg_conf >= pos_conf and neg_conf > 0:
|
|
missing.append(pos)
|
|
else:
|
|
unknown.append(pos)
|
|
return wearing, missing, unknown
|