app.py

import streamlit as st
from pubmed_rag import search_pubmed, fetch_pubmed_abstracts, summarize_text
from image_pipeline import analyze_medical_image
from models import query_medical_text

st.set_page_config(page_title="Advanced Medical AI", layout="wide")


def main():
    st.title("Advanced Medical AI")
    st.sidebar.title("Features")
    task = st.sidebar.selectbox("Choose a task:", ["PubMed Q&A", "Medical Image Analysis"])

    if task == "PubMed Q&A":
        # PubMed Q&A Section
        st.subheader("PubMed Question Answering")
        query = st.text_input("Enter your medical question:", "What are the latest treatments for diabetes?")
        max_results = st.slider("Number of PubMed articles to retrieve:", 1, 10, 5)

        if st.button("Run Query"):
            with st.spinner("Searching PubMed..."):
                pmids = search_pubmed(query, max_results)
                if not pmids:
                    st.error("No results found. Try another query.")
                    return

            with st.spinner("Fetching and summarizing abstracts..."):
                abstracts = fetch_pubmed_abstracts(pmids)
                summaries = {pmid: summarize_text(abstract) for pmid, abstract in abstracts.items()}

            st.subheader("PubMed Summaries")
            for pmid, summary in summaries.items():
                st.write(f"**PMID {pmid}**: {summary}")

            with st.spinner("Querying the medical reasoning model..."):
                answer = query_medical_text(query)
            st.subheader("AI-Powered Answer")
            st.write(answer)

    elif task == "Medical Image Analysis":
        # Medical Image Analysis Section
        st.subheader("Medical Image Analysis")
        uploaded_file = st.file_uploader("Upload a medical image (PNG/JPG):", type=["png", "jpg", "jpeg"])

        if uploaded_file:
            # Display the uploaded image
            st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
            
            # Analyze the medical image
            with st.spinner("Analyzing image..."):
                result = analyze_medical_image(uploaded_file)
            
            # Display the result
            st.subheader("Diagnostic Insight")
            st.write(result)


if __name__ == "__main__":
    main()