learn.py

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()