Add prescription app files

app.py

@@ -1,3 +1,208 @@
+"""
+app.py — Prescription Digitization Tool
+Vision-first pipeline: tries llava/moondream vision model first (most accurate),
+falls back to OCR + Mistral text mode if no vision model available.
+"""
+
+import io, os
+import pandas as pd
 import streamlit as st
+from PIL import Image
+
+import db
+from preprocess import preprocess_image
+from ocr_engine import run_ocr
+from llm_extractor import extract_from_image, extract_fields, DEFAULT_MODEL, DEFAULT_HOST
+
+st.set_page_config(page_title="Prescription Digitization", page_icon="🏥", layout="wide")
+db.init_db()
+
+for k, v in {
+    "ocr_text": "", "extraction_result": None,
+    "ocr_engine_used": "both", "uploaded_filename": "",
+    "image_bytes": None,
+}.items():
+    if k not in st.session_state:
+        st.session_state[k] = v
+
 st.title("🏥 Prescription Digitization")
-st.write("App is running successfully!")
+st.caption("Fully offline · Patient data never leaves this machine")
+
+tab_upload, tab_records, tab_export = st.tabs(
+    ["📤 Upload & Extract", "📋 Saved Records", "⬇️ Export CSV"]
+)
+
+with tab_upload:
+    uploaded = st.file_uploader(
+        "Upload prescription image",
+        type=["jpg", "jpeg", "png"],
+        label_visibility="collapsed",
+        accept_multiple_files=False,
+    )
+
+    if uploaded:
+        raw_bytes = uploaded.read()
+        st.session_state.image_bytes    = raw_bytes
+        st.session_state.uploaded_filename = uploaded.name
+        image = Image.open(io.BytesIO(raw_bytes))
+
+        col_orig, col_clean = st.columns(2)
+        with col_orig:
+            st.image(image, caption="Original", width=450)
+        with col_clean:
+            cleaned = preprocess_image(image, do_deskew=True)
+            st.image(cleaned, caption="Preprocessed", width=450)
+
+        if st.button("🔍 Extract Information", type="primary", use_container_width=True):
+
+            # ── Try vision model first ────────────────────────────────
+            with st.spinner("Checking for vision model (llava/moondream)…"):
+                vision_result = extract_from_image(raw_bytes, host=DEFAULT_HOST)
+
+            if vision_result.success:
+                st.session_state.extraction_result = vision_result
+                st.session_state.ocr_engine_used   = vision_result.mode
+                n = len(vision_result.records)
+                st.success(
+                    f"✨ Vision model extracted **{n} visit record{'s' if n>1 else ''}** "
+                    f"— review below before saving."
+                )
+
+            elif vision_result.error == "NO_VISION_MODEL":
+                # ── Fall back to OCR + Mistral ────────────────────────
+                st.info("No vision model found — using OCR + Mistral (slower, less accurate).\n\n"
+                        "For better results: `ollama pull llava`")
+
+                with st.spinner("Running OCR (Donut + TrOCR)…"):
+                    try:
+                        cleaned_bytes = io.BytesIO()
+                        cleaned.save(cleaned_bytes, format="JPEG")
+                        ocr_result = run_ocr(cleaned, engine="both")
+                        st.session_state.ocr_text        = ocr_result.raw_text
+                        st.session_state.ocr_engine_used = ocr_result.engine
+                    except Exception as e:
+                        st.error(f"OCR failed: {e}")
+                        st.stop()
+
+                with st.spinner(f"Structuring with Ollama ({DEFAULT_MODEL})…"):
+                    text_result = extract_fields(
+                        st.session_state.ocr_text,
+                        model=DEFAULT_MODEL, host=DEFAULT_HOST
+                    )
+
+                if text_result.success:
+                    st.session_state.extraction_result = text_result
+                    n = len(text_result.records)
+                    st.success(f"Found **{n} visit record{'s' if n>1 else ''}** — review below.")
+                elif text_result.error == "OLLAMA_NOT_RUNNING":
+                    st.error("Ollama is not running.")
+                    st.code("ollama serve")
+                    st.session_state.extraction_result = None
+                else:
+                    st.error(f"Extraction error: {text_result.error}")
+                    st.info("Try clicking Extract Information again.")
+                    st.session_state.extraction_result = None
+
+            else:
+                st.error(f"Vision extraction error: {vision_result.error}")
+                st.session_state.extraction_result = None
+
+    if st.session_state.ocr_text:
+        with st.expander("Raw OCR output", expanded=False):
+            st.code(st.session_state.ocr_text, language=None)
+
+    # ── Review form ───────────────────────────────────────────────────
+    result = st.session_state.extraction_result
+    if result and result.records:
+        st.divider()
+        st.subheader("✏️ Review extracted information")
+
+        meta = result.meta
+        if meta:
+            st.caption(f"Mode: {result.mode}")
+            mc = st.columns(min(len(meta), 4))
+            for col, (k, v) in zip(mc, meta.items()):
+                col.text_input(k.replace("_"," ").capitalize(), value=str(v), disabled=True)
+
+        for i, rec in enumerate(result.records):
+            label = f"Visit {i+1}"
+            if rec.get("visit_date"):
+                label += f"  —  {rec['visit_date']}"
+
+            with st.expander(label, expanded=True):
+                with st.form(f"form_{i}"):
+                    edited_fields = {}
+                    field_items   = list(rec["fields"].items())
+
+                    for j in range(0, len(field_items), 3):
+                        chunk = field_items[j:j+3]
+                        cols  = st.columns(len(chunk))
+                        for col, (key, val) in zip(cols, chunk):
+                            edited_fields[key] = col.text_input(
+                                key.replace("_"," ").capitalize(),
+                                value=str(val) if val else ""
+                            )
+
+                    meds = rec.get("medications", [])
+                    if meds:
+                        st.markdown("**Medications**")
+                        med_df = pd.DataFrame(meds)
+                        for c in ["drug_name","dosage","frequency","route"]:
+                            if c not in med_df.columns: med_df[c] = ""
+                        edited_meds_df = st.data_editor(
+                            med_df[["drug_name","dosage","frequency","route"]],
+                            num_rows="dynamic", use_container_width=True, key=f"meds_{i}"
+                        )
+                        final_meds = edited_meds_df.fillna("").to_dict(orient="records")
+                    else:
+                        final_meds = []
+                        st.caption("No medications detected — add manually if needed.")
+
+                    if st.form_submit_button(
+                        f"✅ Save Visit {i+1} to Database",
+                        type="primary", use_container_width=True
+                    ):
+                        clean_fields = {k: v for k, v in edited_fields.items() if v.strip()}
+                        pid = db.save_record(
+                            visit_date=rec.get("visit_date"),
+                            fields=clean_fields,
+                            medications=final_meds,
+                            meta=meta,
+                            ocr_engine=st.session_state.ocr_engine_used,
+                            source_filename=st.session_state.uploaded_filename,
+                        )
+                        st.success(f"Visit {i+1} saved as record #{pid}.")
+
+with tab_records:
+    records = db.fetch_prescriptions()
+    if not records:
+        st.info("No records saved yet.")
+    else:
+        st.dataframe(pd.DataFrame(records), use_container_width=True, hide_index=True)
+        st.subheader("Medications")
+        sel_id = st.selectbox("Select record", [r["id"] for r in records],
+                              format_func=lambda x: f"Record #{x}")
+        if sel_id:
+            meds = db.fetch_medications(sel_id)
+            if meds:
+                st.dataframe(pd.DataFrame(meds), use_container_width=True, hide_index=True)
+            else:
+                st.caption("No medications for this record.")
+        st.divider()
+        del_id = st.number_input("Delete record by ID", min_value=0, step=1, value=0)
+        if st.button("🗑️ Delete record") and del_id > 0:
+            db.delete_prescription(int(del_id))
+            st.success(f"Deleted #{del_id}.")
+            st.rerun()
+
+with tab_export:
+    flat = db.fetch_all_flat()
+    if not flat:
+        st.info("No data to export yet.")
+    else:
+        df = pd.DataFrame(flat)
+        st.dataframe(df, use_container_width=True, hide_index=True)
+        csv = df.to_csv(index=False).encode("utf-8")
+        st.download_button("⬇️ Download CSV", data=csv,
+                           file_name="prescriptions_export.csv", mime="text/csv",
+                           use_container_width=True)