update OCR

app.py

@@ -1,150 +1,182 @@
-# app.py — Gradio demo
-# Đặt ở root project: engineering-drawing-ai/app.py
+import os
+# Tắt các cảnh báo hệ thống của Ultralytics và Paddle trên Hugging Face
+os.environ['YOLO_CONFIG_DIR'] = '/tmp'
+os.environ["PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK"] = "True"
 
-import os, sys, json, tempfile
 import gradio as gr
+import torch
 import cv2
+import json
 import numpy as np
+from ultralytics import RTDETR
+from paddleocr import PaddleOCR
 from PIL import Image
 
-# Auto-download weights từ HuggingFace Hub nếu chưa có
-CHECKPOINT = "best.pt"
-HF_REPO    = "phamha/drawing-model-weights"  # ← sửa sau
-
-def ensure_weights():
-    if not os.path.exists(CHECKPOINT):
-        print("[INFO] Downloading model weights from HuggingFace...")
+# --- Fix lỗi PyTorch load weights ---
+_orig_torch_load = torch.load
+def _patched_load(*args, **kwargs):
+    kwargs.setdefault("weights_only", False)
+    return _orig_torch_load(*args, **kwargs)
+torch.load = _patched_load
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+CLASS_NAMES   = ["note", "part-drawing", "table"]
+CLASS_DISPLAY = {"note": "Note", "part-drawing": "PartDrawing", "table": "Table"}
+COLORS        = {"note": (0,165,255), "part-drawing": (0,200,0), "table": (0,0,220)}
+
+_det_model  = None
+_paddle_ocr = None
+
+def get_det_model():
+    global _det_model
+    if _det_model is None:
+        # Tự động load file best.pt đang có sẵn trong cùng thư mục
+        _det_model = RTDETR('best.pt')
+    return _det_model
+
+def get_ocr_reader():
+    global _paddle_ocr
+    if _paddle_ocr is None:
+        # lang='vi' hỗ trợ cả tiếng Việt, tiếng Anh và số cực tốt
+        _paddle_ocr = PaddleOCR(use_angle_cls=False, lang='vi', show_log=False)
+    return _paddle_ocr
+
+def _group_rows(items):
+    """Thuật toán gom cụm các text cùng 1 hàng cho Table"""
+    if not items:
+        return []
+    items = sorted(items, key=lambda x: x["y"])
+    y_vals = [it["y"] for it in items]
+    if len(y_vals) > 1:
+        gaps   = [y_vals[i+1]-y_vals[i] for i in range(len(y_vals)-1)]
+        thresh = max(8, (sum(gaps)/len(gaps)) * 0.6)
+    else:
+        thresh = 12
+        
+    rows, cur = [], [items[0]]
+    for item in items[1:]:
+        if item["y"] - cur[-1]["y"] < thresh:
+            cur.append(item)
+        else:
+            cur.sort(key=lambda x: x["x"])
+            rows.append([i["text"] for i in cur])
+            cur = [item]
+    cur.sort(key=lambda x: x["x"])
+    rows.append([i["text"] for i in cur])
+    return rows
+
+def process_pipeline(input_image):
+    if input_image is None:
+        return None, "Vui lòng tải ảnh lên.", ""
+
+    # Chuyển PIL Image sang BGR Numpy cho OpenCV
+    img_bgr = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
+    img_draw = img_bgr.copy()
+    
+    # 1. Detection
+    model = get_det_model()
+    results = model(img_bgr, imgsz=1024, conf=0.3, iou=0.5, device=DEVICE, verbose=False)
+    
+    ocr_reader = get_ocr_reader()
+    objects = []
+    ocr_display = ""
+    
+    # 2. Xử lý từng bounding box
+    for i, box in enumerate(results[0].boxes):
+        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
+        cls_idx  = int(box.cls[0])
+        conf_val = round(float(box.conf[0]), 4)
+        
+        # Mapping label 
         try:
-            from huggingface_hub import hf_hub_download
-            hf_hub_download(
-                repo_id=HF_REPO,
-                filename="best.pt",
-                local_dir=".",
-                local_dir_use_symlinks=False,
-            )
-            print("[INFO] Weights downloaded.")
-        except Exception as e:
-            print(f"[ERROR] Cannot download weights: {e}")
-            raise
-
-ensure_weights()
-
-# Import pipeline SAU KHI đảm bảo weights tồn tại
-sys.path.insert(0, ".")
-from src.inference import run_pipeline
-
-
-# ── Xử lý ảnh ──────────────────────────────────────────────
-def process(image: Image.Image):
-    if image is None:
-        return None, "{}", "Chưa có ảnh."
-
-    tmp_dir  = tempfile.mkdtemp()
-    tmp_path = os.path.join(tmp_dir, "input.jpg")
-    image.save(tmp_path, quality=95)
-
-    try:
-        result, vis_path = run_pipeline(
-            image_path  = tmp_path,
-            output_dir  = tmp_dir,
-            checkpoint  = CHECKPOINT,
-            conf_thresh = 0.3,
-        )
-    except Exception as e:
-        import traceback
-        return None, "{}", f"Lỗi pipeline:\n{traceback.format_exc()}"
-
-    # Ảnh kết quả
-    vis_bgr = cv2.imread(vis_path)
-    vis_rgb = cv2.cvtColor(vis_bgr, cv2.COLOR_BGR2RGB)
-
-    # JSON sạch (bỏ crop_path)
-    clean_objs = []
-    for obj in result["objects"]:
-        clean_objs.append({
-            "id":          obj["id"],
-            "class":       obj["class"],
-            "confidence":  obj["confidence"],
-            "bbox":        obj["bbox"],
-            "ocr_content": obj["ocr_content"],
+            cls_raw = CLASS_NAMES[cls_idx]
+            cls_show = CLASS_DISPLAY[cls_raw]
+        except:
+            cls_raw = model.names[cls_idx].lower()
+            cls_show = cls_raw.capitalize()
+
+        pad = 5
+        crop = img_bgr[max(0, y1-pad):min(img_bgr.shape[0], y2+pad),
+                       max(0, x1-pad):min(img_bgr.shape[1], x2+pad)]
+        
+        ocr_content = ""
+        # Thực hiện OCR
+        if cls_raw in ["note", "table"]:
+            # PaddleOCR chạy trực tiếp trên ảnh màu gốc
+            ocr_res = ocr_reader.ocr(crop, cls=False)
+            
+            if ocr_res and ocr_res[0]:
+                if cls_raw == "note":
+                    lines = [line[1][0] for line in ocr_res[0]]
+                    ocr_content = "\n".join(lines)
+                    ocr_display += f"### Note {i+1}:\n{ocr_content}\n\n"
+                    
+                elif cls_raw == "table":
+                    items = []
+                    for line in ocr_res[0]:
+                        box_pts = line[0]
+                        text = line[1][0]
+                        center_x = (box_pts[0][0] + box_pts[1][0]) / 2
+                        center_y = (box_pts[0][1] + box_pts[2][1]) / 2
+                        items.append({"text": text, "x": center_x, "y": center_y})
+                        
+                    rows = _group_rows(items)
+                    ocr_content = "\n".join(" | ".join(r) for r in rows)
+                    ocr_display += f"### Table {i+1}:\n```text\n{ocr_content}\n```\n\n"
+
+        # Lưu dữ liệu JSON
+        objects.append({
+            "id": i+1,
+            "class": cls_show,
+            "confidence": conf_val,
+            "bbox": {"x1": x1, "y1": y1, "x2": x2, "y2": y2},
+            "ocr_content": ocr_content if ocr_content else None
         })
-    json_str = json.dumps(
-        {"image": result["image"], "objects": clean_objs},
-        ensure_ascii=False, indent=2
-    )
 
-    # OCR text đẹp
-    ocr_parts = []
-    for obj in result["objects"]:
-        content = obj.get("ocr_content")
-        if not content:
-            continue
-        if isinstance(content, dict):        # Table
-            content = content.get("text", "")
-        if not content.strip():
-            continue
-        sep = "─" * 46
-        ocr_parts.append(
-            f"{sep}\n"
-            f"[{obj['class']} #{obj['id']}]  conf={obj['confidence']}\n"
-            f"{sep}\n{content}"
-        )
-    ocr_text = "\n\n".join(ocr_parts) or "Không phát hiện Note / Table."
-
-    return vis_rgb, json_str, ocr_text
-
-
-# ── Gradio UI ───────────────────────────────────────────────
-with gr.Blocks(title="Engineering Drawing Analyzer", theme=gr.themes.Soft()) as demo:
-
-    gr.Markdown("""
-    # 🔧 Engineering Drawing Analyzer
-    **Tự động phát hiện và trích xuất văn bản từ bản vẽ kỹ thuật**
-    
-    Hỗ trợ 3 loại vùng:
-    - 🟢 **PartDrawing** — vùng bản vẽ chi tiết
-    - 🟠 **Note** — ghi chú, chú thích
-    - 🔴 **Table** — bảng dữ liệu kỹ thuật
-    """)
+        # Vẽ hình
+        color = COLORS.get(cls_raw, (0, 255, 0))
+        cv2.rectangle(img_draw, (x1, y1), (x2, y2), color, 3)
+        label = f"{cls_show} {conf_val:.2f}"
+        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
+        cv2.rectangle(img_draw, (x1, y1 - th - 10), (x1 + tw + 8, y1), color, -1)
+        cv2.putText(img_draw, label, (x1 + 4, y1 - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
 
-    with gr.Row():
-        with gr.Column(scale=1):
-            inp = gr.Image(type="pil", label="📁 Upload bản vẽ kỹ thuật")
-            btn = gr.Button("🔍 Detect & OCR", variant="primary", size="lg")
+    # Convert ảnh kết quả trả về Gradio
+    final_img = Image.fromarray(cv2.cvtColor(img_draw, cv2.COLOR_BGR2RGB))
+    
+    json_result = {
+        "image": "uploaded_image.jpg",
+        "objects": objects
+    }
+    json_str = json.dumps(json_result, ensure_ascii=False, indent=2)
+    
+    if not ocr_display:
+        ocr_display = "Không tìm thấy text nào trong các vùng Note/Table."
 
-        with gr.Column(scale=1):
-            out_img = gr.Image(label="✅ Kết quả detection")
+    return final_img, json_str, ocr_display
 
+# --- GIAO DIỆN GRADIO ---
+with gr.Blocks(title="Engineering Drawing OCR") as demo:
+    gr.Markdown("# 🏭 Trích xuất dữ liệu Bản vẽ Kỹ thuật (Detection & OCR)")
+    gr.Markdown("**Mô hình Detection:** RT-DETR | **Mô hình OCR:** PaddleOCR")
+    
     with gr.Row():
         with gr.Column(scale=1):
-            out_json = gr.Code(
-                language="json",
-                label="📋 JSON output",
-                lines=25,
-            )
+            img_input = gr.Image(type="pil", label="Tải ảnh bản vẽ lên")
+            btn_submit = gr.Button("🔍 Phân tích", variant="primary")
         with gr.Column(scale=1):
-            out_ocr = gr.Textbox(
-                label="📝 OCR content (Note & Table)",
-                lines=25,
-                max_lines=60,
-            )
-
-    btn.click(
-        fn      = process,
-        inputs  = [inp],
-        outputs = [out_img, out_json, out_ocr],
+            img_output = gr.Image(type="pil", label="Kết quả Detection")
+            
+    with gr.Row():
+        json_output = gr.Code(language="json", label="Dữ liệu JSON")
+        ocr_output = gr.Markdown(label="Kết quả OCR")
+        
+    btn_submit.click(
+        fn=process_pipeline,
+        inputs=[img_input],
+        outputs=[img_output, json_output, ocr_output]
     )
 
-    gr.Markdown("""
-    ---
-    **Model:** RT-DETR-L fine-tuned | **OCR:** EasyOCR (vi+en) + PaddleOCR fallback  
-    **mAP50:** 0.942 | **Dataset:** Engineering drawings (Vietnamese technical)
-    """)
-
-
 if __name__ == "__main__":
-    demo.launch(
-        server_name = "0.0.0.0",
-        server_port = 7860,
-        share       = False,   # đổi True nếu muốn link public tạm
-    )
+    demo.launch()