[ocr]: add overlays for color and state estimation

anpr_ocr/logger.py

@@ -62,24 +62,23 @@ def estimate_vehicle_color(frame_bgr: np.ndarray, bounding_box: Any) -> str:
     median_value = float(np.median(val))
     median_saturation = float(np.median(sat))
     if median_value < 55:
-        return "Black"
-    if median_saturation < 35:
-        brightness = median_value
-        return "White" if brightness > 185 else "Silver/Gray"
-
-    dominant_hue = float(np.median(hue))
-    if dominant_hue < 10 or dominant_hue >= 170:
-        color = "Red"
-    elif dominant_hue < 25:
-        color = "Orange/Brown"
-    elif dominant_hue < 35:
-        color = "Yellow/Gold"
-    elif dominant_hue < 85:
-        color = "Green"
-    elif dominant_hue < 135:
-        color = "Blue"
+        color = "Black"
+    elif median_saturation < 35:
+        color = "White" if median_value > 185 else "Silver/Gray"
     else:
-        color = "Purple"
+        dominant_hue = float(np.median(hue))
+        if dominant_hue < 10 or dominant_hue >= 170:
+            color = "Red"
+        elif dominant_hue < 25:
+            color = "Orange/Brown"
+        elif dominant_hue < 35:
+            color = "Yellow/Gold"
+        elif dominant_hue < 85:
+            color = "Green"
+        elif dominant_hue < 135:
+            color = "Blue"
+        else:
+            color = "Purple"
     return color
 
 

anpr_ocr/utils.py

@@ -136,7 +136,6 @@ def enhance_plate_image(
         "CH",
         "CG",
         "DD",
-        "DH",
         "DL",
         "DN",
         "GA",
@@ -180,8 +179,7 @@ def enhance_plate_image(
     "BR": "Bihar",
     "CH": "Chandigarh",
     "CG": "Chhattisgarh",
-    "DD": "Daman and Diu",
-    "DH": "Dadra and Nagar Haveli",
+    "DD": "Dadra and Nagar Haveli and Daman and Diu",
     "DL": "Delhi",
     "DN": "Dadra and Nagar Haveli and Daman and Diu",
     "GA": "Goa",

anpr_ocr/video_demo.py

@@ -15,6 +15,8 @@
 
 from anpr_ocr import ALPR, PlateLogger
 from anpr_ocr.alpr import SUPPORTED_VIDEO_EXTS
+from anpr_ocr.logger import estimate_vehicle_color
+from anpr_ocr.utils import get_state_name
 
 # pylint: disable=too-many-branches, too-many-statements, import-outside-toplevel
 # ruff: noqa: PLR0912, PLR0915, PLC0415, E501, ARG001
@@ -222,7 +224,7 @@ def _play_video_live(
     from anpr_ocr.utils import PlateTracker
 
     tracker = PlateTracker(max_unseen_frames=frame_skip * 5, window_size=5)
-    active_plates: list[tuple[Any, str, float]] = []
+    active_plates: list[tuple[Any, str, float, str, str]] = []
     frame_idx = 0
     t_start = time.perf_counter()
 
@@ -254,10 +256,19 @@ def _play_video_live(
                 frame_dets.append((r.detection.bounding_box, r.ocr.text.strip(), conf))
 
             tracked = tracker.update(frame_dets, frame_idx)
-            active_plates = [(box, txt, conf) for box, txt, conf, _ in tracked]
+            active_plates = [
+                (
+                    box,
+                    txt,
+                    conf,
+                    get_state_name(txt) or "Other / International",
+                    estimate_vehicle_color(frame, box),
+                )
+                for box, txt, conf, _ in tracked
+            ]
 
             if logger is not None:
-                for box, txt, conf in active_plates:
+                for box, txt, conf, _, _ in active_plates:
                     logger.observe(
                         plate_text=txt,
                         confidence=conf,
@@ -269,9 +280,9 @@ def _play_video_live(
 
         # Render annotations
         display = frame.copy()
-        for b, txt, c in active_plates:
+        for b, txt, c, state, color in active_plates:
             cv2.rectangle(display, (b.x1, b.y1), (b.x2, b.y2), (36, 255, 12), 2)
-            lbl = f"{txt} {c * 100:.0f}%"
+            lbl = f"{txt} {c * 100:.0f}% | {color} | {state}"
             cv2.putText(
                 display,
                 lbl,

test/test_alpr.py

@@ -22,6 +22,7 @@
     pad_bounding_box,
     split_two_row_crop,
     vote_consensus_plate,
+    get_state_name,
 )
 
 ASSETS_DIR = Path(__file__).resolve().parent.parent / "assets"
@@ -211,6 +212,26 @@ def __init__(self, x1: int, y1: int, x2: int, y2: int) -> None:
     assert res2[0][3] == tid
 
 
+def test_indian_state_code_mapping() -> None:
+    assert get_state_name("MH12DE1433") == "Maharashtra"
+    assert get_state_name("DL01AB1234") == "Delhi"
+    assert get_state_name("OD02AB1234") == "Odisha"
+    assert get_state_name("DD01AB1234") == "Dadra and Nagar Haveli and Daman and Diu"
+    assert get_state_name("XX01AB1234") == ""
+
+
+def test_vehicle_color_estimation() -> None:
+    class DummyBox:
+        def __init__(self, x1: int, y1: int, x2: int, y2: int) -> None:
+            self.x1, self.y1, self.x2, self.y2 = x1, y1, x2, y2
+
+    box = DummyBox(40, 40, 80, 60)
+    blue_frame = np.full((120, 120, 3), (255, 0, 0), dtype=np.uint8)
+    gray_frame = np.full((120, 120, 3), (150, 150, 150), dtype=np.uint8)
+    assert estimate_vehicle_color(blue_frame, box) == "Blue"
+    assert estimate_vehicle_color(gray_frame, box) == "Silver/Gray"
+
+
 def test_plate_logger_and_exports(tmp_path: Path) -> None:
     class DummyBox:
         def __init__(self, x1: int, y1: int, x2: int, y2: int) -> None: