| 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%}") |
