From f30e1b7fd89cb5ee4fe3a68825a8bef67a74ef3b Mon Sep 17 00:00:00 2001 From: Hardik Kumar Date: Mon, 7 Sep 2026 19:53:16 +0530 Subject: [PATCH] [ocr]: add overlays for color and state estimation - Added real-time vehicle color estimation. - Added Indian state/UT name resolution Signed-off-by: Hardik Kumar --- anpr_ocr/logger.py | 49 ++++++++++++++++++++++++++++++++++++++++++ anpr_ocr/utils.py | 4 +--- anpr_ocr/video_demo.py | 21 +++++++++++++----- test/test_alpr.py | 23 +++++++++++++++++++- 4 files changed, 88 insertions(+), 9 deletions(-) diff --git a/anpr_ocr/logger.py b/anpr_ocr/logger.py index 70f5742..65cced0 100644 --- a/anpr_ocr/logger.py +++ b/anpr_ocr/logger.py @@ -33,6 +33,55 @@ def is_similar_plate(p1: str, p2: str, threshold: float = 0.75) -> bool: return SequenceMatcher(None, p1, p2).ratio() >= threshold +def estimate_vehicle_color(frame_bgr: np.ndarray, bounding_box: Any) -> str: + """Estimate the dominant visible vehicle color around a plate detection.""" + height, width = frame_bgr.shape[:2] + box_width = max(1, bounding_box.x2 - bounding_box.x1) + box_height = max(1, bounding_box.y2 - bounding_box.y1) + x1 = max(0, bounding_box.x1 - box_width) + y1 = max(0, bounding_box.y1 - box_height * 2) + x2 = min(width, bounding_box.x2 + box_width) + y2 = min(height, bounding_box.y2 + box_height * 2) + region = frame_bgr[y1:y2, x1:x2] + if region.size == 0: + return "Unknown" + + hsv = cv2.cvtColor(region, cv2.COLOR_BGR2HSV) + pixels = hsv.reshape(-1, 3) + saturation = pixels[:, 1] + value = pixels[:, 2] + valid = pixels[(value > 25) & (saturation > 20)] + if len(valid) < 20: + valid = pixels[value > 25] + if len(valid) < 20: + return "Unknown" + + hue = valid[:, 0].astype(np.float64) + sat = valid[:, 1].astype(np.float64) + val = valid[:, 2].astype(np.float64) + median_value = float(np.median(val)) + median_saturation = float(np.median(sat)) + if median_value < 55: + color = "Black" + elif median_saturation < 35: + color = "White" if median_value > 185 else "Silver/Gray" + else: + 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 + + @dataclass class VehicleRecord: """A finalized log record for a unique vehicle at its peak recognition score.""" diff --git a/anpr_ocr/utils.py b/anpr_ocr/utils.py index 12332da..ba1cc5b 100644 --- a/anpr_ocr/utils.py +++ b/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", diff --git a/anpr_ocr/video_demo.py b/anpr_ocr/video_demo.py index 6837e55..b2f8b36 100644 --- a/anpr_ocr/video_demo.py +++ b/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, diff --git a/test/test_alpr.py b/test/test_alpr.py index f4c264e..2d0768e 100644 --- a/test/test_alpr.py +++ b/test/test_alpr.py @@ -12,7 +12,7 @@ from anpr_ocr.alpr import ALPR from anpr_ocr.cli import _create_main_parser -from anpr_ocr.logger import PlateLogger +from anpr_ocr.logger import PlateLogger, estimate_vehicle_color from anpr_ocr.utils import ( PlateTracker, disambiguate_plate, @@ -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: