File size: 4,372 Bytes
b42373a 56f5eaf b42373a | 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 | import cv2
import os
import json
import numpy as np
from preprocess import preprocess_chassis
CONFIG_PATH = "config.json"
_ocr = None
def load_config():
if os.path.exists(CONFIG_PATH):
with open(CONFIG_PATH) as f:
return json.load(f)
return {
"confusion_map": {},
"error_correction": {"max_errors": 2, "window_size": 3}
}
def get_ocr():
global _ocr
if _ocr is None:
from paddleocr import PaddleOCR
_ocr = PaddleOCR(use_angle_cls=True, lang='en')
return _ocr
def ocr_image(img):
ocr = get_ocr()
result = ocr.ocr(img, cls=True)
if not result or not result[0]:
return "", 0.0
texts = [line[1][0] for line in result[0]]
confs = [line[1][1] for line in result[0]]
full_text = "".join(texts).upper()
full_text = "".join(c for c in full_text if c.isalnum())
avg_conf = sum(confs) / len(confs) if confs else 0.0
return full_text, avg_conf
def read_chassis(image_path, save_comparison=False):
variations = preprocess_chassis(image_path, save_comparison=save_comparison)
best_text, best_conf, best_score = "", 0.0, -1
for var in variations:
text, conf = ocr_image(var)
score = conf * max(len(text), 1)
if score > best_score:
best_text, best_conf, best_score = text, conf, score
return best_text, best_conf
def can_substitute(got, want, confusion_map):
return want in confusion_map.get(got, []) or got in confusion_map.get(want, [])
def apply_substitutions(ocr_text, expected_text, confusion_map, max_errors):
if len(ocr_text) != len(expected_text):
return ocr_text, False
diffs = [(i, ocr_text[i], expected_text[i])
for i in range(len(ocr_text)) if ocr_text[i] != expected_text[i]]
if len(diffs) > max_errors:
return ocr_text, False
corrected = list(ocr_text)
for i, got, want in diffs:
if can_substitute(got, want, confusion_map):
corrected[i] = want
else:
return ocr_text, False
return "".join(corrected), True
def best_window_match(ocr_text, expected_text, window_size):
exp_len = len(expected_text)
best, best_diffs = None, exp_len + 1
for start in range(max(0, len(ocr_text) - exp_len) + 1):
candidate = ocr_text[start:start + exp_len]
if len(candidate) != exp_len:
continue
diffs = sum(1 for a, b in zip(candidate, expected_text) if a != b)
if diffs < best_diffs:
best_diffs = diffs
best = (candidate, diffs)
if best and best[1] <= window_size:
return best
return None
def postprocess_with_hint(ocr_text, expected_text):
config = load_config()
confusion_map = config.get("confusion_map", {})
ec = config.get("error_correction", {"max_errors": 2, "window_size": 3})
max_errors = ec["max_errors"]
window_size = ec["window_size"]
if not ocr_text:
return ocr_text, False
if ocr_text == expected_text:
return ocr_text, True
if expected_text in ocr_text:
return expected_text, True
if len(ocr_text) == len(expected_text):
corrected, fixed = apply_substitutions(
ocr_text, expected_text, confusion_map, max_errors)
if fixed:
return corrected, True
if abs(len(ocr_text) - len(expected_text)) <= window_size:
match = best_window_match(ocr_text, expected_text, window_size)
if match:
candidate, diffs = match
if diffs == 0:
return candidate, True
corrected, fixed = apply_substitutions(
candidate, expected_text, confusion_map, max_errors)
if fixed:
return corrected, True
if len(ocr_text) < len(expected_text):
suffix = expected_text[-len(ocr_text):]
corrected, fixed = apply_substitutions(
ocr_text, suffix, confusion_map, max_errors=1)
if fixed or ocr_text == suffix:
return expected_text, True
return ocr_text, False
if __name__ == "__main__":
import sys
if len(sys.argv) < 2:
print("Usage: python ocr.py <chassis_image_path>")
else:
text, conf = read_chassis(sys.argv[1], save_comparison=True)
print(f"Result : {text}")
print(f"Confidence : {conf:.2%}") |