import streamlit as st
from transformers import pipeline
import pytesseract
import cv2
import numpy as np
from PIL import Image

# Load the model
model = pipeline("text2text-generation", model="google/paligemma-3b-mix-224")

# Function to extract text from image using OCR
def extract_text_from_image(image_file):
    image = Image.open(image_file)
    img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
    text = pytesseract.image_to_string(img_cv)
    return text

# Function to get AI interpretation of the prescription
def interpret_prescription(text):
    response = model(text)
    return response[0]['generated_text'].strip()

# Set Streamlit page configuration
st.set_page_config(
    page_title="Prescription Reader",
    page_icon="💊",
    layout="centered",
)

# Header
st.title("Doctor's Prescription Reader 💊")

# Upload prescription image
uploaded_file = st.file_uploader("Upload Prescription Image", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Display uploaded image
    st.image(uploaded_file, caption="Uploaded Prescription", use_column_width=True)

    with st.spinner("Extracting text from prescription..."):
        # Extract text from image using OCR
        extracted_text = extract_text_from_image(uploaded_file)
        st.subheader("Extracted Text from Prescription:")
        st.text(extracted_text)

        if extracted_text:
            # Interpret extracted text using the model
            with st.spinner("Interpreting the prescription..."):
                ai_response = interpret_prescription(extracted_text)
                st.subheader("AI Interpretation:")
                st.text(ai_response)