Spaces:
Running on Zero
Running on Zero
File size: 6,758 Bytes
de9997b 9fc77da de9997b 9fc77da de9997b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 | import numpy as np
import torch
from typing import List
from src.utils.numeric import extract_numeric_value
from src.utils.postprocessing import hungarian_match
class NumericOCR:
def __init__(
self,
language: str = 'en',
rec_model: str = 'latin_PP-OCRv5_mobile_rec'
):
"""
Initializes PaddleOCR.
:param language: PaddleOCR language code for latin-script recognition
:param rec_model: PaddleOCR recognition model name
"""
self.language = language
self.rec_model = rec_model
self.available = False
try:
from paddleocr import PaddleOCR
# enable_mkldnn=False: oneDNN's CPU kernels crash with
# 'RuntimeError: could not execute a primitive' on some image
# crops (PaddlePaddle/PaddleOCR#14103); disabling MKLDNN makes
# CPU inference slower but robust.
self._engine = PaddleOCR(
use_angle_cls=True,
lang=self.language,
text_recognition_model_name=self.rec_model,
enable_mkldnn=False,
show_log=False
)
self.available = True
except ImportError:
print('PaddleOCR not available, numeric OCR will be disabled')
@staticmethod
def _normalize_polygon(
polygon_xy: List[List[float]],
width: int,
height: int
) -> List[List[float]]:
"""
Normalizes polygon points from [x, y] pixels to [y, x] in [0, 1].
:param polygon_xy: polygon points in pixel coordinates [x, y]
:param width: image width
:param height: image height
:return: normalized polygon points [[y, x], ...]
"""
norm = []
w = float(width)
h = float(height)
for x, y in polygon_xy:
x = float(x)
y = float(y)
x_n = min(max(x / w, 0.0), 1.0)
y_n = min(max(y / h, 0.0), 1.0)
norm.append([y_n, x_n])
return norm
@staticmethod
def _ransac_predict_bar_values(
p_bars: torch.Tensor,
matched_tick_centers: torch.Tensor,
matched_values: torch.Tensor
) -> torch.Tensor:
"""
RANSAC fitting and bar value prediction from bar y-coordinates.
:param p_bars: predicted bars [y, x, conf], shape [N, 3]
:param matched_tick_centers: matched tick positions [y, x], shape [M, 2]
:param matched_values: numeric values at matched tick positions, shape [M]
:return: predicted bars with values [B, 3] as [y, x, value], sorted left -> right
"""
device = p_bars.device
if matched_tick_centers.size(0) < 2:
return torch.empty((0, 3), device=device)
# Sort matched ticks bottom -> top before fitting
order = torch.argsort(matched_tick_centers[:, 0], descending=True)
matched_tick_centers = matched_tick_centers[order]
matched_values = matched_values[order]
from sklearn.linear_model import RANSACRegressor
ys_np = matched_tick_centers[:, 0].detach().cpu().numpy().reshape(-1, 1)
values_np = matched_values.detach().cpu().numpy()
ransac = RANSACRegressor(random_state=0)
ransac.fit(ys_np, values_np)
slope = torch.tensor(float(ransac.estimator_.coef_[0]), device=device)
intercept = torch.tensor(float(ransac.estimator_.intercept_), device=device)
# Predict values at bar y-coordinates and sort left -> right
bar_centers = p_bars[:, :2]
pred_values = slope * bar_centers[:, 0] + intercept
order = torch.argsort(bar_centers[:, 1])
sorted_centers = bar_centers[order]
sorted_values = pred_values[order].unsqueeze(1)
return torch.cat((sorted_centers, sorted_values), dim=1)
def infer_bar_values(
self,
bars: torch.Tensor,
ticks: torch.Tensor,
image: np.ndarray,
dist_thresh: float,
gt_tick_yxv: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Infers bar values via OCR and oracle (GT tick labels).
:param bars: predicted bars [y, x, conf], shape [N, 3]
:param ticks: predicted ticks [y, x, conf], shape [M, 3]
:param image: full-resolution image array, shape [H, W, 3], dtype uint8
:param dist_thresh: max. spatial distance for tick matching
:param gt_tick_yxv: ground-truth ticks [y, x, value], shape [K, 3]
:return: tuple containing:
- ocr bars, [B, 3] as [y, x, value], sorted left -> right
- oracle bars, [B, 3] as [y, x, value], sorted left -> right
"""
device = bars.device
empty = torch.empty((0, 3), device=device)
tick_centers = ticks[:, :2]
# Oracle path
if bars.size(0) == 0 or ticks.size(0) == 0 or gt_tick_yxv.size(0) == 0:
oracle = empty
else:
tick_idx, gt_idx = hungarian_match(tick_centers, gt_tick_yxv[:, :2], dist_thresh)
if tick_idx.numel() == 0:
oracle = empty
else:
oracle = self._ransac_predict_bar_values(
bars,
tick_centers[tick_idx],
gt_tick_yxv[:, 2][gt_idx]
)
# OCR path
if bars.size(0) == 0 or ticks.size(0) == 0 or not self.available:
return empty, oracle
h, w = image.shape[:2]
raw = self._engine.ocr(image, cls=True)
pages = raw or []
centers = []
values = []
for page in pages:
for det in page:
polygon = det[0]
text_score = det[1]
text = str(text_score[0]).strip()
value = extract_numeric_value(text)
if value is None:
continue
poly_xy = [[float(x), float(y)] for x, y in polygon]
poly_yx = self._normalize_polygon(poly_xy, w, h)
center = np.mean(np.array(poly_yx, dtype=np.float32), axis=0).tolist()
centers.append([float(center[0]), float(center[1])])
values.append(float(value))
if len(values) == 0:
return empty, oracle
text_centers = torch.tensor(centers, device=device)
text_values = torch.tensor(values, device=device)
tick_idx, text_idx = hungarian_match(tick_centers, text_centers, dist_thresh)
if tick_idx.numel() == 0:
return empty, oracle
ocr = self._ransac_predict_bar_values(
bars,
tick_centers[tick_idx],
text_values[text_idx]
)
return ocr, oracle
|