feat: use easyocr and enhance llm prompts

app.py

@@ -2,23 +2,25 @@ import json
 import re
 from typing import Any, Dict, List, Optional, Sequence
 
+import easyocr
 import gradio as gr
+import numpy as np
 import torch
 from PIL import Image, ImageDraw
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 # --- OCR pipeline ---------------------------------------------------------
 # Use a high-capacity OCR model for better accuracy on prescription labels.
-OCR_MODEL_ID = "microsoft/trocr-large-printed"
+OCR_LANGS = ["ko", "en"]
 LLM_MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
 
 
 def _load_ocr():
-    device = 0 if torch.cuda.is_available() else -1
-    return pipeline("image-to-text", model=OCR_MODEL_ID, device=device)
+    use_gpu = torch.cuda.is_available()
+    return easyocr.Reader(OCR_LANGS, gpu=use_gpu)
 
 
-ocr = _load_ocr()
+ocr_reader = _load_ocr()
 
 
 def _load_llm():
@@ -153,7 +155,23 @@ def parse_fields(raw: str) -> Dict[str, Any]:
 
 
 def ocr_and_parse(image: Image.Image) -> Dict[str, Any]:
-    raw_text = ocr(image)[0]["generated_text"]
+    np_img = np.array(image.convert("RGB"))
+    results = ocr_reader.readtext(np_img, detail=1, paragraph=False)
+
+    segments: List[Dict[str, Any]] = []
+    lines: List[str] = []
+    for bbox, text, confidence in results:
+        cleaned = text.strip()
+        if not cleaned:
+            continue
+        lines.append(cleaned)
+        segments.append({
+            "text": cleaned,
+            "confidence": float(confidence),
+            "bbox": bbox,
+        })
+
+    raw_text = "\n".join(lines)
     fields = parse_fields(raw_text)
 
     warnings: List[str] = []
@@ -162,7 +180,12 @@ def ocr_and_parse(image: Image.Image) -> Dict[str, Any]:
     if not fields["times_per_day"]:
         warnings.append("1일 횟수를 찾지 못했습니다 (예: 1일 3회).")
 
-    return {"raw_text": raw_text, "fields": fields, "warnings": warnings}
+    return {
+        "raw_text": raw_text,
+        "fields": fields,
+        "warnings": warnings,
+        "segments": segments,
+    }
 
 
 def render_card(fields: Dict[str, Any]) -> Image.Image:
@@ -270,8 +293,18 @@ def generate_llm_explanations(output: Dict[str, Any]) -> str:
         "당신은 약사 선생님입니다. 어려운 의학 용어를 쓰지 말고, 중학생도 이해할 수 있는 말투로 친절하게 설명하세요."
     )
     user_prompt = (
-        "다음은 약봉투 OCR 결과입니다. 약 이름과 용량 정보를 참고해 각 약의 역할을 쉽게 설명하고, 언제 복용하면 좋은지 예시, 주의사항을 bullet로 정리해 주세요.\n"
-        f"약 목록:\n{context}\n\nOCR 원문:\n{raw_text}\n\n출력 형식:\n- 약 이름: ...\n  - 한 줄 설명\n  - 예시 상황\n  - 주의할 점\n마지막에는 의료진 복약 지시를 반드시 따라야 한다는 문장을 덧붙여 주세요."
+        "다음은 약봉투에서 OCR로 추출한 전체 텍스트입니다. 약 이름과 복용 지시를 기반으로 각 약의 정보를 아주 쉽게 정리해 주세요.\n"
+        "요구 사항:\n"
+        "1. 각 약마다 아래 항목을 bullet 형식으로 작성합니다.\n"
+        "   - 약 이름: (가능하면 한글/영문 병기)\n"
+        "   - 어떤 약인지 한 줄 설명\n"
+        "   - 복용 예시: 언제, 어떤 상황에서 복용하면 좋은지 예시\n"
+        "   - 복용 방법 예시: 1회 용량/하루 횟수가 있다면 언급\n"
+        "   - 부작용 또는 주의사항: 흔한 부작용, 피해야 할 행동\n"
+        "2. 어려운 의학 용어는 피하고, 중학생도 이해할 수 있는 말투로 작성합니다.\n"
+        "3. 약 이름을 확실히 모르면 ‘이름 미확인’이라고 쓰고, 약사에게 확인하라고 안내합니다.\n"
+        "4. 마지막 문단에 반드시 ‘실제 복약은 의사·약사의 지시를 따르세요’ 문장을 포함하세요.\n"
+        f"\n약 목록(추출 요약):\n{context}\n\nOCR 원문 전체:\n{raw_text}\n"
     )
 
     messages = [

requirements.txt

@@ -3,3 +3,6 @@ torch
 gradio
 Pillow
 sentencepiece
+easyocr
+opencv-python-headless
+numpy