Rewrite Gradio app with proper parsing

app.py

@@ -1,113 +1,158 @@
-# app.py (HF Space)
-
-
-# dose_per_intake: e.g., "1회 1정", "1정", "5 mL"
-m_dose = re.search(r"(1회\s*)?(\d+)\s*([가-힣a-zA-Z]+|mL|ml|mg)", t)
-dose_per_intake = None
-if m_dose:
-dose_per_intake = f"{m_dose.group(2)} {m_dose.group(3)}"
-
-
-# drug name (heuristic): token before mg/mL or first uppercase-like word
-m_drug = re.search(r"([가-힣A-Za-z]+)\s*(\d+\s*(mg|mL|ml))", t)
-drug_name = m_drug.group(1) if m_drug else None
-
-
-return {
-"drug_name": drug_name,
-"dose_per_intake": dose_per_intake,
-"times_per_day": times_per_day,
-"time_slots": time_slots or None,
-}
-
-
-
-
-def ocr_and_parse(img) -> Dict[str, Any]:
-"""Run OCR then parse fields with basic validation."""
-# OCR output is a list of dicts with 'generated_text'
-raw = ocr(img)[0]["generated_text"]
-fields = parse_fields(raw)
-
-
-# basic validation messages
-warn = []
-if not fields["drug_name"]:
-warn.append("약 이름 인식이 불확실합니다.")
-if not fields["times_per_day"]:
-warn.append("1일 횟수를 찾지 못했습니다 (예: 1일 3회).")
-
-
-return {"raw_text": raw, "fields": fields, "warnings": warn}
+import json
+import re
+from typing import Any, Dict, List, Optional
+
+import gradio as gr
+from PIL import Image, ImageDraw
+from transformers import pipeline
+
+# --- OCR pipeline ---------------------------------------------------------
+# We use a light-weight printed-text OCR model that works well for receipts/labels.
+ocr = pipeline("image-to-text", model="microsoft/trocr-base-printed")
+
+# Korean keywords describing time slots on prescription labels.
+TIME_KEYWORDS = [
+    "아침",
+    "점심",
+    "저녁",
+    "취침",
+    "자기",
+    "식전",
+    "식후",
+    "식간",
+    "기상",
+]
 
 
+def _extract_time_slots(text: str) -> List[str]:
+    slots = []
+    for kw in TIME_KEYWORDS:
+        if kw in text:
+            slots.append(kw)
+    # Also capture explicit times like 08:00 혹은 8시
+    for match in re.findall(r"(\d{1,2}[:시]\d{0,2})", text):
+        norm = match.replace("시", ":")
+        if norm.endswith(":"):
+            norm += "00"
+        if norm not in slots:
+            slots.append(norm)
+    return slots
+
+
+def parse_fields(raw: str) -> Dict[str, Any]:
+    """Extract drug name and dosage information from OCR text."""
+    text = raw.replace("\n", " ")
+    text = re.sub(r"\s+", " ", text)
+
+    # 1) 약 이름: 단어 + 용량 패턴 주변에서 찾기
+    drug_name: Optional[str] = None
+    drug_match = re.search(r"([가-힣A-Za-z]+)\s*(\d+)\s*(mg|mL|ML|정)", text)
+    if drug_match:
+        drug_name = drug_match.group(1)
+    else:
+        fallback = re.search(r"([가-힣A-Za-z]{2,})", text)
+        drug_name = fallback.group(1) if fallback else None
+
+    # 2) 1회 용량: "1회 1정", "1정", "5 mL" 등
+    dose_per_intake: Optional[str] = None
+    dose_match = re.search(r"(1회\s*)?(\d+[\./]?\d*)\s*([가-힣A-Za-z]+|mL|ml|mg|정)", text)
+    if dose_match:
+        dose_per_intake = f"{dose_match.group(2)} {dose_match.group(3)}".strip()
+
+    # 3) 1일 복용 횟수: "1일 3회", "하루 2회"
+    times_per_day: Optional[int] = None
+    times_match = re.search(r"(?:1일|하루)\s*(\d+)\s*회", text)
+    if times_match:
+        times_per_day = int(times_match.group(1))
+
+    # 4) 시간대 키워드/시각 추출
+    time_slots = _extract_time_slots(text)
+
+    return {
+        "drug_name": drug_name,
+        "dose_per_intake": dose_per_intake,
+        "times_per_day": times_per_day,
+        "time_slots": time_slots or None,
+    }
+
+
+def ocr_and_parse(image: Image.Image) -> Dict[str, Any]:
+    raw_text = ocr(image)[0]["generated_text"]
+    fields = parse_fields(raw_text)
+
+    warnings: List[str] = []
+    if not fields["drug_name"]:
+        warnings.append("약 이름 인식이 불확실합니다.")
+    if not fields["times_per_day"]:
+        warnings.append("1일 횟수를 찾지 못했습니다 (예: 1일 3회).")
+
+    return {"raw_text": raw_text, "fields": fields, "warnings": warnings}
 
 
 def render_card(fields: Dict[str, Any]) -> Image.Image:
-"""Draw a simple schedule card (PNG) from parsed fields."""
-W, H = 720, 400
-img = Image.new("RGB", (W, H), "white")
-d = ImageDraw.Draw(img)
-title = "오늘 복용 일정"
-d.rectangle((0, 0, W, 60), fill=(230, 240, 255))
-d.text((24, 18), title, fill=(0, 0, 0))
+    width, height = 720, 400
+    img = Image.new("RGB", (width, height), "white")
+    draw = ImageDraw.Draw(img)
 
+    header_text = "오늘 복용 일정"
+    draw.rectangle((0, 0, width, 60), fill=(230, 240, 255))
+    draw.text((24, 18), header_text, fill=(0, 0, 0))
 
-y = 90
-def line(label, value):
-nonlocal y
-d.text((24, y), f"{label}", fill=(60, 60, 60))
-d.text((180, y), f": {value if value else '-'}", fill=(0, 0, 0))
-y += 34
+    y = 90
 
+    def add_line(label: str, value: Optional[str]):
+        nonlocal y
+        draw.text((24, y), label, fill=(60, 60, 60))
+        display = value if value else "-"
+        draw.text((180, y), f": {display}", fill=(0, 0, 0))
+        y += 34
 
-line("약 이름", fields.get("drug_name"))
-line("1회 용량", fields.get("dose_per_intake"))
-line("1일 횟수", fields.get("times_per_day"))
-slots = ", ".join(fields.get("time_slots") or [])
-line("시간대", slots if slots else None)
+    add_line("약 이름", fields.get("drug_name"))
+    add_line("1회 용량", fields.get("dose_per_intake"))
+    add_line("1일 횟수", str(fields.get("times_per_day") or ""))
 
+    slots = fields.get("time_slots") or []
+    add_line("시간대", ", ".join(slots) if slots else None)
 
-d.text((24, H-60), "※ 의료진 처방 우선, 본 앱은 안내용입니다.", fill=(120, 120, 120))
-return img
+    footer = "※ 의료진 처방이 우선이며, 본 앱은 참고용입니다."
+    draw.text((24, height - 60), footer, fill=(120, 120, 120))
+    return img
 
 
+def to_csv_row(output: Dict[str, Any]) -> str:
+    fields = output["fields"]
+    row = [
+        fields.get("drug_name") or "",
+        fields.get("dose_per_intake") or "",
+        str(fields.get("times_per_day") or ""),
+        ";".join(fields.get("time_slots") or []),
+    ]
+    return ",".join(row)
 
 
-def run_pipeline(image):
-if image is None:
-return "이미지를 업로드하세요.", None, None
-out = ocr_and_parse(image)
-card = render_card(out["fields"])
-csv_row = to_csv_row(out)
-return json.dumps(out, ensure_ascii=False, indent=2), card, csv_row
+def run_pipeline(image: Optional[Image.Image]):
+    if image is None:
+        return "이미지를 업로드하세요.", None, None
 
-
-
-
-def to_csv_row(out: Dict[str, Any]) -> str:
-f = out["fields"]
-row = [
-f.get("drug_name") or "",
-f.get("dose_per_intake") or "",
-str(f.get("times_per_day") or ""),
-";".join(f.get("time_slots") or []),
-]
-return ",".join(row)
+    output = ocr_and_parse(image)
+    card = render_card(output["fields"])
+    csv_row = to_csv_row(output)
+    json_text = json.dumps(output, ensure_ascii=False, indent=2)
+    return json_text, card, csv_row
 
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-gr.Markdown("# MedCard-KR · 약봉투 OCR → 복용 일정 카드")
-with gr.Row():
-with gr.Column():
-img_in = gr.Image(type="pil", label="약 봉투/라벨 사진")
-btn = gr.Button("인식 & 카드 생성", variant="primary")
-csv = gr.Textbox(label="CSV(약명,1회용량,1일횟수,시간대)")
-with gr.Column():
-json_out = gr.Code(label="인식 결과(JSON)")
-card = gr.Image(type="pil", label="일정 카드(미리보기)")
-btn.click(run_pipeline, inputs=img_in, outputs=[json_out, card, csv])
+    gr.Markdown("# MedCard-KR · 약봉투 OCR → 복용 일정 카드")
+    with gr.Row():
+        with gr.Column():
+            img_in = gr.Image(type="pil", label="약 봉투/라벨 사진")
+            btn = gr.Button("인식 & 카드 생성", variant="primary")
+            csv_box = gr.Textbox(label="CSV(약명,1회용량,1일횟수,시간대)")
+        with gr.Column():
+            json_out = gr.Code(label="인식 결과(JSON)")
+            card_out = gr.Image(type="pil", label="일정 카드(미리보기)")
+    btn.click(run_pipeline, inputs=img_in, outputs=[json_out, card_out, csv_box])
 
 
 if __name__ == "__main__":
-demo.queue().launch()
+    demo.queue().launch()