HungLuong10/microservice

1from pix2tex.cli import LatexOCR
2from PIL import Image
3import hashlib
4import os
5import time
6import json
7
8# ========= CONFIG =========
9IMAGE_DIR = "data"
10OUTPUT_JSON = "ocr_latex.json"
11
12# ========= OCR INIT (1 LẦN) =========
13ocr = LatexOCR()
14
15# ========= CACHE =========
16ocr_cache = {}  # sha1 -> latex
17
18
19def sha1_file(path: str) -> str:
20    h = hashlib.sha1()
21    with open(path, "rb") as f:
22        for chunk in iter(lambda: f.read(8192), b""):
23            h.update(chunk)
24    return h.hexdigest()
25
26
27def ocr_with_cache(image_path: str):
28    sha1 = sha1_file(image_path)
29
30    if sha1 in ocr_cache:
31        return {
32            "sha1": sha1,
33            "latex": ocr_cache[sha1],
34            "cache": True,
35            "time_sec": 0.0,
36        }
37
38    start = time.perf_counter()
39    img = Image.open(image_path)
40    latex = ocr(img)
41    elapsed = time.perf_counter() - start
42
43    ocr_cache[sha1] = latex
44
45    return {
46        "sha1": sha1,
47        "latex": latex,
48        "cache": False,
49        "time_sec": round(elapsed, 4),
50    }
51
52
53# ========= BENCHMARK =========
54results = []
55total_start = time.perf_counter()
56
57png_files = sorted(
58    f for f in os.listdir(IMAGE_DIR)
59    if f.lower().endswith(".png")
60)
61
62print(f"Found {len(png_files)} PNG files")
63
64for idx, fname in enumerate(png_files, 1):
65    path = os.path.join(IMAGE_DIR, fname)
66
67    print(f"[{idx}/{len(png_files)}] OCR {fname}")
68    r = ocr_with_cache(path)
69
70    results.append({
71        "file": fname,
72        "sha1": r["sha1"],
73        "latex": r["latex"],
74        "cache": r["cache"],
75        "time_sec": r["time_sec"],
76    })
77
78total_time = time.perf_counter() - total_start
79
80summary = {
81    "total_images": len(png_files),
82    "unique_images": len(ocr_cache),
83    "total_time_sec": round(total_time, 3),
84    "avg_time_per_image_sec": round(
85        total_time / max(len(png_files), 1), 4
86    ),
87}
88
89output = {
90    "summary": summary,
91    "results": results,
92}
93
94# ========= SAVE JSON =========
95with open(OUTPUT_JSON, "w", encoding="utf-8") as f:
96    json.dump(output, f, ensure_ascii=False, indent=2)
97
98print("\n===== BENCHMARK SUMMARY =====")
99for k, v in summary.items():
100    print(f"{k}: {v}")
101
102print(f"\nSaved results to {OUTPUT_JSON}")
103