| import os |
| import json |
| import cv2 |
| import numpy as np |
| from collections import defaultdict |
| from barcode_scanner import scan_all_barcodes |
| from preprocess import preprocess_chassis |
|
|
| BARCODE_DIR = "images/barcode" |
| CHASSIS_DIR = "images/chassis" |
| CONFIG_PATH = "config.json" |
|
|
|
|
| def run_ocr_ensemble(image_path, ocr): |
| variations = preprocess_chassis(image_path) |
| best_text, best_score = "", -1 |
| for var in variations: |
| result = ocr.ocr(var, cls=True) |
| if not result or not result[0]: |
| continue |
| texts = [line[1][0] for line in result[0]] |
| confs = [line[1][1] for line in result[0]] |
| text = "".join(texts).upper() |
| text = "".join(c for c in text if c.isalnum()) |
| conf = sum(confs) / len(confs) if confs else 0.0 |
| score = conf * max(len(text), 1) |
| if score > best_score: |
| best_text, best_score = text, score |
| return best_text |
|
|
|
|
| def best_alignment(got, expected): |
| exp_len = len(expected) |
| if len(got) == exp_len: |
| return got |
| best_start, best_diffs = 0, exp_len + 1 |
| for start in range(max(0, len(got) - exp_len) + 1): |
| cand = got[start:start + exp_len] |
| if len(cand) != exp_len: |
| continue |
| diffs = sum(1 for a, b in zip(cand, expected) if a != b) |
| if diffs < best_diffs: |
| best_diffs = diffs |
| best_start = start |
| return got[best_start:best_start + exp_len] |
|
|
|
|
| def learn_confusion_map(ocr_results, ground_truths): |
| counts = defaultdict(lambda: defaultdict(int)) |
| for key, expected in ground_truths.items(): |
| got = ocr_results.get(key, "") |
| if not got or got == expected: |
| continue |
| aligned = best_alignment(got, expected) |
| if len(aligned) != len(expected): |
| continue |
| for g, e in zip(aligned, expected): |
| if g != e: |
| counts[g][e] += 1 |
|
|
| confusion_map = {} |
| print("\n Learned confusions:") |
| for char in sorted(counts.keys()): |
| wants = sorted(counts[char], key=lambda w: counts[char][w], reverse=True) |
| confusion_map[char] = wants |
| print(f" '{char}' -> {wants} (counts: {dict(counts[char])})") |
|
|
| return confusion_map |
|
|
|
|
| def can_fix(got_str, expected_str, confusion_map, max_errors): |
| if len(got_str) != len(expected_str): |
| return False |
| diffs = [(g, e) for g, e in zip(got_str, expected_str) if g != e] |
| if len(diffs) > max_errors: |
| return False |
| return all( |
| e in confusion_map.get(g, []) or g in confusion_map.get(e, []) |
| for g, e in diffs |
| ) |
|
|
|
|
| def learn_thresholds(ocr_results, ground_truths, confusion_map): |
| best_correct, best_config = 0, {"max_errors": 2, "window_size": 3} |
|
|
| for max_err in [1, 2, 3, 4]: |
| for win in [2, 3, 4, 5]: |
| correct = 0 |
| for key, expected in ground_truths.items(): |
| got = ocr_results.get(key, "") |
| if not got: |
| continue |
| if got == expected or expected in got: |
| correct += 1 |
| continue |
| if abs(len(got) - len(expected)) <= win: |
| aligned = best_alignment(got, expected) |
| if can_fix(aligned, expected, confusion_map, max_err): |
| correct += 1 |
| elif len(got) < len(expected): |
| suffix = expected[-len(got):] |
| if can_fix(got, suffix, confusion_map, 1): |
| correct += 1 |
|
|
| if correct > best_correct: |
| best_correct = correct |
| best_config = {"max_errors": max_err, "window_size": win} |
|
|
| print(f"\n Best thresholds: {best_config} " |
| f"(estimated correct: {best_correct}/{len(ground_truths)})") |
| return best_config |
|
|
|
|
| def main(): |
| print("=" * 60) |
| print("LEARNING FROM DATA") |
| print("=" * 60) |
|
|
| print("\n[1/3] Scanning barcodes for ground truth...") |
| ground_truths = scan_all_barcodes(BARCODE_DIR) |
| ground_truths = {k: v for k, v in ground_truths.items() if v} |
| print(f" Got {len(ground_truths)} ground truth labels") |
|
|
| print("\n[2/3] Running ensemble OCR on all chassis images...") |
| from paddleocr import PaddleOCR |
| ocr = PaddleOCR(use_angle_cls=True, lang='en', |
| use_gpu=False, show_log=False) |
|
|
| chassis_files = sorted([ |
| f for f in os.listdir(CHASSIS_DIR) |
| if f.lower().endswith(('.jpg', '.jpeg', '.png')) |
| ]) |
|
|
| ocr_results = {} |
| for fname in chassis_files: |
| key = os.path.splitext(fname)[0] |
| path = os.path.join(CHASSIS_DIR, fname) |
| text = run_ocr_ensemble(path, ocr) |
| ocr_results[key] = text |
| expected = ground_truths.get(key, "???") |
| match = "[OK]" if text == expected else "[--]" |
| print(f" {match} {key}: got='{text}' expected='{expected}'") |
|
|
| print("\n[3/3] Learning confusion map and thresholds...") |
| confusion_map = learn_confusion_map(ocr_results, ground_truths) |
| thresholds = learn_thresholds(ocr_results, ground_truths, confusion_map) |
|
|
| existing = {} |
| if os.path.exists(CONFIG_PATH): |
| with open(CONFIG_PATH) as f: |
| existing = json.load(f) |
|
|
| config = { |
| "confusion_map": confusion_map, |
| "preprocessing": existing.get("preprocessing", { |
| "clahe_clip": 3.0, |
| "clahe_grid": 8, |
| "bilateral_d": 9, |
| "bilateral_sigma": 75, |
| "adaptive_blocksize": 21, |
| "adaptive_c": 8, |
| "padding": 20 |
| }), |
| "error_correction": thresholds |
| } |
|
|
| with open(CONFIG_PATH, "w") as f: |
| json.dump(config, f, indent=2) |
|
|
| print(f"\n[DONE] Config saved -> {CONFIG_PATH}") |
| print("Now run: python evaluate.py") |
| print("=" * 60) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
