update OCR

src/inference.py

@@ -31,34 +31,32 @@ def get_det_model(checkpoint="best.pt"):
     return _det_model
 
 def get_ocr_reader(backend="paddle"):
-    """Khởi tạo OCR engine, ưu tiên PaddleOCR, fallback EasyOCR"""
     global _ocr_reader
     if _ocr_reader is not None:
         return _ocr_reader
-
     if backend == "paddle":
         try:
             from paddleocr import PaddleOCR
-            print("[INFO] Initializing PaddleOCR (lang: vi, en)...")
+            print("[INFO] Initializing PaddleOCR (optimized for table & vi)...")
             _ocr_reader = PaddleOCR(
-                lang='vi',                     # tiếng Việt + tiếng Anh
-                use_angle_cls=True,            # tự động xoay ảnh
+                lang='vi',
+                use_angle_cls=True,
                 use_gpu=(DEVICE == "cuda"),
                 show_log=False,
-                det_db_thresh=0.3,
-                det_db_box_thresh=0.5,
-                rec_algorithm='SVTR_LCNet'     # mạnh cho chữ in
+                det_db_thresh=0.2,        # giảm ngưỡng để phát hiện nhiều chữ hơn
+                det_db_box_thresh=0.3,
+                det_db_unclip_ratio=1.5,   # mở rộng box
+                rec_algorithm='SVTR_LCNet',
+                rec_image_shape="3,48,320", # hỗ trợ chữ dài
+                max_text_length=50
             )
             return _ocr_reader
-        except ImportError:
-            print("[WARN] PaddleOCR not installed, falling back to EasyOCR.")
         except Exception as e:
-            print(f"[WARN] PaddleOCR init failed: {e}, fallback to EasyOCR.")
-
-    # Fallback to EasyOCR
-    import easyocr
-    print("[INFO] Loading EasyOCR (vi, en)...")
-    _ocr_reader = easyocr.Reader(["vi", "en"], gpu=(DEVICE == "cuda"), verbose=False)
+            print(f"[WARN] PaddleOCR failed: {e}, fallback to EasyOCR")
+            backend = "easyocr"
+    if backend == "easyocr":
+        import easyocr
+        _ocr_reader = easyocr.Reader(["vi","en"], gpu=(DEVICE=="cuda"), verbose=False)
     return _ocr_reader
 
 def preprocess_image(img_bgr, ocr_type="note"):
@@ -183,70 +181,149 @@ def group_rows(items, vertical_thresh_ratio=0.6):
     return [[it["text"] for it in row] for row in rows]
 
 def ocr_note(img_path, backend="paddle"):
-    """OCR cho vùng Note, trả về chuỗi văn bản."""
     img = cv2.imread(img_path)
     if img is None:
         return ""
+    
+    # Tiền xử lý mạnh: tăng độ phân giải, lọc nhiễu, làm sắc nét
+    h, w = img.shape[:2]
+    if w < 1000:
+        scale = 1200 / w
+        img = cv2.resize(img, (int(w*scale), int(h*scale)), interpolation=cv2.INTER_CUBIC)
+    
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    # Sử dụng CLAHE mạnh hơn
+    clahe = cv2.createCLAHE(clipLimit=3.5, tileGridSize=(4,4))
+    gray = clahe.apply(gray)
+    # Khử nhiễu nhưng giữ cạnh
+    gray = cv2.bilateralFilter(gray, 9, 75, 75)
+    # Làm sắc nét
+    kernel_sharpen = np.array([[-1,-1,-1],[-1,9,-1],[-1,-1,-1]])
+    gray = cv2.filter2D(gray, -1, kernel_sharpen)
+    # Nhị phân hóa thích ứng
+    gray = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)
+    img_proc = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
+    
+    reader = get_ocr_reader(backend)
+    if hasattr(reader, 'ocr'):
+        result = reader.ocr(img_proc, cls=True)
+        if result and result[0]:
+            texts = [line[1][0].strip() for line in result[0] if line[1][1] > 0.3]
+            return "\n".join(texts)
+    else:
+        results = reader.readtext(img_proc, detail=1, paragraph=True)
+        texts = [text for (_, text, conf) in results if conf > 0.3]
+        return "\n".join(texts)
+    return ""
 
-    items = ocr_with_backend(img, backend, ocr_type="note")
-    if not items:
-        # Thử lại với preprocessing khác (bỏ sharpen, chỉ CLAHE)
-        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
-        gray = clahe.apply(gray)
-        img2 = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
-        items = ocr_with_backend(img2, backend, ocr_type="note")
-        if not items:
-            return ""
-
-    # Sắp xếp theo y rồi x để tạo đoạn văn bản
-    items_sorted = sorted(items, key=lambda x: (x["y"], x["x"]))
-    lines = []
-    current_line = []
-    y_thresh = 12  # ngưỡng dòng
-    for i, it in enumerate(items_sorted):
-        if i == 0:
-            current_line.append(it["text"])
-        else:
-            if abs(it["y"] - items_sorted[i-1]["y"]) < y_thresh:
-                current_line.append(it["text"])
-            else:
-                lines.append(" ".join(current_line))
-                current_line = [it["text"]]
-    if current_line:
-        lines.append(" ".join(current_line))
-
-    # Post-processing: loại bỏ ký tự lạ, chuẩn hóa khoảng trắng
-    clean_lines = []
-    for line in lines:
-        line = re.sub(r'[^\w\s\.\,\-\/\(\)]', '', line)  # giữ chữ, số, dấu c��u cơ bản
-        line = re.sub(r'\s+', ' ', line).strip()
-        if len(line) > 1:
-            clean_lines.append(line)
-    return "\n".join(clean_lines)
+import cv2
+import numpy as np
+
+def detect_table_structure(img_bgr):
+    """
+    Phát hiện các ô trong bảng dựa trên đường kẻ.
+    Trả về: danh sách các ô (x1,y1,x2,y2) theo thứ tự từ trái sang phải, trên xuống dưới.
+    """
+    gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
+    # Nhị phân hóa
+    _, binary = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY_INV)
+    # Phát hiện đường kẻ ngang
+    horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (40, 1))
+    horizontal_lines = cv2.morphologyEx(binary, cv2.MORPH_OPEN, horizontal_kernel, iterations=2)
+    # Phát hiện đường kẻ dọc
+    vertical_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 40))
+    vertical_lines = cv2.morphologyEx(binary, cv2.MORPH_OPEN, vertical_kernel, iterations=2)
+    # Kết hợp
+    table_structure = cv2.add(horizontal_lines, vertical_lines)
+    # Tìm contours của các ô (các vùng kín)
+    contours, _ = cv2.findContours(table_structure, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+    cells = []
+    for cnt in contours:
+        x, y, w, h = cv2.boundingRect(cnt)
+        if w > 10 and h > 10:  # lọc ô nhỏ do nhiễu
+            cells.append((x, y, x+w, y+h))
+    # Loại bỏ trùng lặp và sắp xếp theo y trước, x sau
+    cells = sorted(set(cells), key=lambda r: (r[1], r[0]))
+    return cells
+
+def ocr_cell(img_cell, backend="paddle"):
+    """OCR riêng một ô, trả về text đã làm sạch."""
+    if img_cell.size == 0:
+        return ""
+    # Tiền xử lý mạnh cho ô nhỏ
+    gray = cv2.cvtColor(img_cell, cv2.COLOR_BGR2GRAY)
+    # Resize nếu quá nhỏ
+    h, w = gray.shape
+    if w < 50 or h < 30:
+        scale = max(2.0, 100 / w)
+        gray = cv2.resize(gray, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
+    # CLAHE + threshold
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(4,4))
+    gray = clahe.apply(gray)
+    _, gray = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+    # Đưa về BGR cho OCR engine
+    img_proc = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
+    
+    reader = get_ocr_reader(backend)
+    if hasattr(reader, 'ocr'):
+        result = reader.ocr(img_proc, cls=True)
+        if result and result[0]:
+            texts = [line[1][0].strip() for line in result[0] if line[1][1] > 0.3]
+            return " ".join(texts)
+    else:
+        results = reader.readtext(img_proc, detail=1, paragraph=False)
+        texts = [text for (_, text, conf) in results if conf > 0.2]
+        return " ".join(texts)
+    return ""
 
 def ocr_table(img_path, backend="paddle"):
-    """OCR cho vùng Table, trả về dict rows và text."""
+    """Phiên bản cải tiến: phát hiện ô, OCR từng ô, ghép theo hàng."""
     img = cv2.imread(img_path)
     if img is None:
         return {"rows": [], "text": ""}
-
+    
+    # Phát hiện cấu trúc bảng (các ô)
+    cells = detect_table_structure(img)
+    if not cells:
+        # Fallback: dùng phương pháp cũ (group theo dòng)
+        return ocr_table_fallback(img, backend)
+    
+    # OCR từng ô
+    ocr_results = []
+    for (x1, y1, x2, y2) in cells:
+        # Mở rộng nhẹ để lấy toàn bộ chữ
+        pad = 2
+        y1 = max(0, y1-pad)
+        x1 = max(0, x1-pad)
+        y2 = min(img.shape[0], y2+pad)
+        x2 = min(img.shape[1], x2+pad)
+        cell_img = img[y1:y2, x1:x2]
+        text = ocr_cell(cell_img, backend)
+        if text:
+            ocr_results.append({
+                "text": text,
+                "x": (x1 + x2) // 2,
+                "y": (y1 + y2) // 2,
+                "box": (x1, y1, x2, y2)
+            })
+    
+    if not ocr_results:
+        return {"rows": [], "text": ""}
+    
+    # Nhóm thành các hàng dựa trên tọa độ y
+    rows = group_rows(ocr_results, vertical_thresh_ratio=0.5)
+    # Mỗi hàng đã được sắp xếp theo x, chỉ lấy text
+    rows_text = [[item["text"] for item in row] for row in rows]
+    return {"rows": rows_text, "text": "\n".join(" | ".join(r) for r in rows_text)}
+
+def ocr_table_fallback(img, backend="paddle"):
+    """Phương pháp cũ dùng grouping đơn giản khi không phát hiện được cấu trúc bảng."""
     items = ocr_with_backend(img, backend, ocr_type="table")
     if not items:
-        # Thử lại với ảnh gốc (không xóa đường kẻ) vì đôi khi đường kẻ giúp định vị ô
-        img2 = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-        img2 = cv2.adaptiveThreshold(img2, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
-                                     cv2.THRESH_BINARY, 15, 5)
-        img2 = cv2.cvtColor(img2, cv2.COLOR_GRAY2BGR)
-        items = ocr_with_backend(img2, backend, ocr_type="table")
-        if not items:
-            return {"rows": [], "text": ""}
-
+        return {"rows": [], "text": ""}
     rows = group_rows(items, vertical_thresh_ratio=0.6)
-    # Chuyển rows thành text (các cột cách nhau bằng ' | ')
-    text_lines = [" | ".join(row) for row in rows if row]
-    return {"rows": rows, "text": "\n".join(text_lines)}
-
+    rows_text = [[it["text"] for it in row] for row in rows]
+    return {"rows": rows_text, "text": "\n".join(" | ".join(r) for r in rows_text)}
 def run_pipeline(image_path, output_dir="outputs",
                  checkpoint="best.pt", conf_thresh=0.3,
                  ocr_backend="paddle"):