import streamlit as st
from PyPDF2 import PdfReader
from transformers import pipeline
from PIL import Image

# Load Hugging Face model
model = pipeline("text-generation", model="Ketak-ZoomRx/Drug_Prompt_Ollama_67k")

def extract_text_from_pdf(pdf_file):
    reader = PdfReader(pdf_file)
    text = ""
    for page in reader.pages:
        text += page.extract_text()
    return text

def get_lab_report_interpretation(text):
    result = model(text)
    return result[0]['label']

def process_image(image_file):
    img = Image.open(image_file)
    return img

# Streamlit UI
st.title("Medical Lab Report Analyzer")

# Upload image and PDF
uploaded_pdf = st.file_uploader("Upload your PDF Medical Lab Report", type="pdf")
uploaded_image = st.file_uploader("Upload Image of Medical Report", type="jpg png jpeg")

if uploaded_pdf is not None:
    st.subheader("Extracted Text from PDF:")
    pdf_text = extract_text_from_pdf(uploaded_pdf)
    st.write(pdf_text)
    
    st.subheader("Lab Report Interpretation:")
    interpretation = get_lab_report_interpretation(pdf_text)
    st.write(interpretation)

elif uploaded_image is not None:
    st.image(process_image(uploaded_image), caption="Uploaded Image", use_column_width=True)
    st.warning("For text extraction from image, please upload a PDF.")
    
else:
    st.info("Please upload a medical PDF or image for analysis.")