Spaces:
Running on Zero
Running on Zero
| """ | |
| utils.py — Reusable Computer Vision and Geometry Utilities for Meter OCR Pipeline | |
| Optimized for Hugging Face Spaces & Standalone Inference. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| from typing import Any | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| def rotate_image(img_bgr: np.ndarray, angle: int) -> np.ndarray: | |
| """ | |
| Rotate BGR image by 90, 180, or 270 degrees clockwise. | |
| Returns original image unchanged if angle == 0. | |
| """ | |
| if angle == 90: | |
| return cv2.rotate(img_bgr, cv2.ROTATE_90_CLOCKWISE) | |
| if angle == 180: | |
| return cv2.rotate(img_bgr, cv2.ROTATE_180) | |
| if angle == 270: | |
| return cv2.rotate(img_bgr, cv2.ROTATE_90_COUNTERCLOCKWISE) | |
| return img_bgr | |
| def apply_prep( | |
| img_bgr: np.ndarray, | |
| prep: str, | |
| clahe_clip: float = 2.0, | |
| clahe_grid: tuple[int, int] = (8, 8), | |
| ) -> np.ndarray: | |
| """ | |
| Apply image contrast enhancement: | |
| - 'clahe': Contrast Limited Adaptive Histogram Equalization on L channel (LAB color space) | |
| - 'histeq': Global Histogram Equalization on Y channel (YCrCb color space) | |
| - 'orig' or unknown: returns original image without modification | |
| """ | |
| if prep == "clahe": | |
| lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB) | |
| l, a, b = cv2.split(lab) | |
| clahe = cv2.createCLAHE(clipLimit=clahe_clip, tileGridSize=clahe_grid) | |
| l_clahe = clahe.apply(l) | |
| return cv2.cvtColor(cv2.merge([l_clahe, a, b]), cv2.COLOR_LAB2BGR) | |
| if prep == "histeq": | |
| ycrcb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2YCrCb) | |
| y, cr, cb = cv2.split(ycrcb) | |
| y_eq = cv2.equalizeHist(y) | |
| return cv2.cvtColor(cv2.merge([y_eq, cr, cb]), cv2.COLOR_YCrCb2BGR) | |
| return img_bgr | |
| def remap_point(x: float, y: float, angle: int, w: int, h: int) -> tuple[float, float]: | |
| """ | |
| Map coordinate (x, y) from a rotated image frame back to original image dimensions (w, h). | |
| """ | |
| if angle == 90: | |
| return y, h - x | |
| if angle == 180: | |
| return w - x, h - y | |
| if angle == 270: | |
| return w - y, x | |
| return x, y | |
| def remap_bbox(bbox: list[float], angle: int, w: int, h: int) -> list[float]: | |
| """ | |
| Map bounding box [x1, y1, x2, y2] from rotated frame back to original image coordinates. | |
| Maintains [min_x, min_y, max_x, max_y] canonical ordering. | |
| """ | |
| x1, y1, x2, y2 = bbox | |
| p1_x, p1_y = remap_point(x1, y1, angle, w, h) | |
| p2_x, p2_y = remap_point(x2, y2, angle, w, h) | |
| return [ | |
| min(p1_x, p2_x), | |
| min(p1_y, p2_y), | |
| max(p1_x, p2_x), | |
| max(p1_y, p2_y), | |
| ] | |
| def inverse_remap_bbox(bbox: list[float], angle: int, w: int, h: int) -> list[float]: | |
| """ | |
| Map bounding box [x1, y1, x2, y2] from original image dimensions to rotated image frame. | |
| Useful for drawing detections directly onto the selected orientation image. | |
| """ | |
| x1, y1, x2, y2 = bbox | |
| if angle == 0: | |
| return bbox | |
| if angle == 90: | |
| return [h - y2, x1, h - y1, x2] | |
| if angle == 180: | |
| return [w - x2, h - y2, w - x1, h - y1] | |
| if angle == 270: | |
| return [y1, w - x2, y2, w - x1] | |
| return bbox | |
| def iou(box1: list[float], box2: list[float]) -> float: | |
| """Calculate Intersection over Union (IoU) between two bounding boxes [x1, y1, x2, y2].""" | |
| x1 = max(box1[0], box2[0]) | |
| y1 = max(box1[1], box2[1]) | |
| x2 = min(box1[2], box2[2]) | |
| y2 = min(box1[3], box2[3]) | |
| intersection = max(0.0, x2 - x1) * max(0.0, y2 - y1) | |
| area1 = max(0.0, box1[2] - box1[0]) * max(0.0, box1[3] - box1[1]) | |
| area2 = max(0.0, box2[2] - box2[0]) * max(0.0, box2[3] - box2[1]) | |
| union = area1 + area2 - intersection | |
| return (intersection / union) if union > 0.0 else 0.0 | |
| def dedup_detections(dets: list[dict[str, Any]], thresh: float = 0.45) -> list[dict[str, Any]]: | |
| """Deduplicate overlapping bounding boxes, preserving higher-confidence candidates.""" | |
| kept: list[dict[str, Any]] = [] | |
| for d in sorted(dets, key=lambda x: x["confidence"], reverse=True): | |
| if not any(iou(d["bbox"], k["bbox"]) > thresh for k in kept): | |
| kept.append(d) | |
| return sorted(kept, key=lambda x: x["center_x"]) | |
| def red_ratio(img_bgr: np.ndarray, bbox: list[float]) -> float: | |
| """Calculate ratio of red pixels within a bounding box (indicating decimal/fraction digits).""" | |
| x1, y1, x2, y2 = [int(v) for v in bbox] | |
| pad_x = max(1, int((x2 - x1) * 0.15)) | |
| pad_y = max(1, int((y2 - y1) * 0.15)) | |
| crop = img_bgr[ | |
| max(0, y1 + pad_y): min(img_bgr.shape[0], y2 - pad_y), | |
| max(0, x1 + pad_x): min(img_bgr.shape[1], x2 - pad_x), | |
| ] | |
| if crop.size == 0: | |
| return 0.0 | |
| hsv = cv2.cvtColor(crop, cv2.COLOR_BGR2HSV) | |
| mask1 = cv2.inRange(hsv, (0, 40, 40), (12, 255, 255)) | |
| mask2 = cv2.inRange(hsv, (165, 40, 40), (180, 255, 255)) | |
| red_mask = mask1 | mask2 | |
| return float(np.mean(red_mask > 0)) | |
| def draw_annotated_image( | |
| image_rgb: np.ndarray, | |
| digits: list[dict[str, Any]], | |
| best_angle: int = 0, | |
| prep: str = "orig", | |
| conf_reliable: float = 0.60, | |
| show_on_rotated: bool = True, | |
| ) -> np.ndarray: | |
| """ | |
| Render bounding boxes, digit labels, and confidence tags on the image. | |
| If show_on_rotated is True, renders on the best-detected rotation and contrast view. | |
| Otherwise renders on the original orientation. | |
| """ | |
| if image_rgb is None or not image_rgb.size: | |
| return image_rgb | |
| h, w = image_rgb.shape[:2] | |
| bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR) | |
| if show_on_rotated and best_angle != 0: | |
| canvas_bgr = rotate_image(bgr, best_angle) | |
| if prep == "clahe": | |
| canvas_bgr = apply_prep(canvas_bgr, "clahe") | |
| elif prep == "histeq": | |
| canvas_bgr = apply_prep(canvas_bgr, "histeq") | |
| else: | |
| canvas_bgr = bgr.copy() | |
| canvas = cv2.cvtColor(canvas_bgr, cv2.COLOR_BGR2RGB) | |
| cur_h, cur_w = canvas.shape[:2] | |
| line_thickness = max(2, int(round(min(cur_w, cur_h) / 320))) | |
| font_scale = max(0.45, min(cur_w, cur_h) / 850.0) | |
| for d in digits: | |
| bbox = d.get("bbox", []) | |
| if len(bbox) != 4: | |
| continue | |
| if show_on_rotated and best_angle != 0: | |
| rx1, ry1, rx2, ry2 = inverse_remap_bbox(bbox, best_angle, w, h) | |
| else: | |
| rx1, ry1, rx2, ry2 = bbox | |
| ix1, iy1 = max(0, int(round(rx1))), max(0, int(round(ry1))) | |
| ix2, iy2 = min(cur_w - 1, int(round(rx2))), min(cur_h - 1, int(round(ry2))) | |
| conf = d.get("confidence", 0.0) | |
| digit_val = d.get("digit", "?") | |
| is_reliable = d.get("reliable", conf >= conf_reliable) | |
| # Colors: Green for reliable (>= 0.60), Amber/Orange for warning (< 0.60) | |
| box_color = (46, 204, 113) if is_reliable else (255, 140, 0) | |
| # Draw bounding box | |
| cv2.rectangle(canvas, (ix1, iy1), (ix2, iy2), box_color, line_thickness) | |
| # Label tag | |
| label = f"{digit_val} ({conf:.2f})" | |
| (tw, th), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 1) | |
| tag_top = max(0, iy1 - th - 6) | |
| tag_bottom = iy1 | |
| tag_right = min(cur_w, ix1 + tw + 8) | |
| # Label background | |
| cv2.rectangle(canvas, (ix1, tag_top), (tag_right, tag_bottom), box_color, -1) | |
| # Text label (black on vibrant tag for clear contrast) | |
| cv2.putText( | |
| canvas, | |
| label, | |
| (ix1 + 4, tag_bottom - 4), | |
| cv2.FONT_HERSHEY_SIMPLEX, | |
| font_scale, | |
| (0, 0, 0), | |
| max(1, int(line_thickness / 2)), | |
| cv2.LINE_AA, | |
| ) | |
| return canvas | |