Update app.py

app.py

@@ -1,136 +1,193 @@
+# --- Docstring ---
+"""
+Streamlit application for Medical Image Analysis using Google Gemini Vision
+and Retrieval-Augmented Generation (RAG) with Chroma DB.
+
+Allows users to upload a medical image (pathology slide, diagram, etc.).
+1. The image is analyzed by Google's Gemini Pro Vision model to generate a
+   textual description of key features.
+2. This description is then used as a query to a Chroma vector database
+   (populated with example medical text snippets) to retrieve relevant
+   information from a simulated knowledge base.
+"""
+
+# --- Imports ---
 import streamlit as st
 import google.generativeai as genai
 import chromadb
 from chromadb.utils import embedding_functions
 from PIL import Image
-import os
 import io
-import time # To create unique IDs for Chroma
+import time # Used for generating unique IDs for Chroma DB demo data
+from typing import Optional, Dict, List, Any # For type hinting
 
 # --- Configuration ---
 try:
-    # Try loading secrets from Hugging Face secrets first
+    # Attempt to load the Google API key from Streamlit secrets
     GOOGLE_API_KEY = st.secrets["GOOGLE_API_KEY"]
     genai.configure(api_key=GOOGLE_API_KEY)
 except KeyError:
-    st.error("GOOGLE_API_KEY not found in Hugging Face secrets!")
+    st.error("❌ GOOGLE_API_KEY not found in Streamlit secrets! Please add it.")
     st.stop()
 except Exception as e:
-    st.error(f"Error configuring Google AI: {e}")
+    st.error(f"❌ Error configuring Google AI SDK: {e}")
     st.stop()
 
 # --- Gemini Model Setup ---
-# Check available models if needed, select the vision model
-# for m in genai.list_models():
-#     if 'generateContent' in m.supported_generation_methods:
-#         print(m.name) # Find the vision model name (e.g., 'gemini-pro-vision')
-
+# Define the specific Gemini model to use (ensure it's a vision-capable model)
 VISION_MODEL_NAME = "gemini-pro-vision"
+
+# Configure generation parameters for the model
+# Lower temperature for more deterministic, factual descriptions
 GENERATION_CONFIG = {
-    "temperature": 0.2, # Lower temp for more factual descriptions
+    "temperature": 0.2,
     "top_p": 0.95,
     "top_k": 40,
     "max_output_tokens": 1024,
 }
-SAFETY_SETTINGS = [ # Adjust safety settings as needed for medical content
+
+# Configure safety settings (adjust thresholds as needed for medical content)
+# Blocking potentially sensitive content might be necessary depending on the images
+SAFETY_SETTINGS = [
     {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
     {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
     {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
     {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
 ]
 
+# Initialize the Gemini Generative Model
 try:
     gemini_model = genai.GenerativeModel(
         model_name=VISION_MODEL_NAME,
         generation_config=GENERATION_CONFIG,
         safety_settings=SAFETY_SETTINGS
     )
+    st.success(f"✅ Initialized Gemini Model: {VISION_MODEL_NAME}")
 except Exception as e:
-    st.error(f"Error initializing Gemini Model ({VISION_MODEL_NAME}): {e}")
+    st.error(f"❌ Error initializing Gemini Model ({VISION_MODEL_NAME}): {e}")
     st.stop()
 
 # --- Chroma DB Setup ---
-# Using persistent storage within the HF Space (data lost if space is wiped)
-# For production, consider a hosted Chroma or other DB solution.
+# Using persistent storage within the Streamlit deployment environment (e.g., HF Space)
+# NOTE: Data will be lost if the persistent storage is wiped or the environment resets.
+# For production, consider a managed Chroma instance or alternative database.
 CHROMA_PATH = "chroma_data"
 COLLECTION_NAME = "medical_docs"
 
-# Use a default sentence transformer embedding function (runs locally on HF space CPU)
-# For better domain adaptation, consider finetuned medical embeddings if possible/available.
-# Make sure the model used here matches the one used when INGESTING data.
+# Define the embedding function.
+# Using a default Sentence Transformer model (runs locally on CPU).
+# IMPORTANT: The embedding model used for querying MUST match the one used
+#            when initially adding data to the collection.
+# For improved performance/relevance on medical text, consider fine-tuned
+# medical domain-specific embedding models if available.
 embedding_func = embedding_functions.DefaultEmbeddingFunction()
 
 try:
+    # Initialize Chroma DB client with persistence
     chroma_client = chromadb.PersistentClient(path=CHROMA_PATH)
-    # Get or create the collection with the specified embedding function
+
+    # Get or create the collection, specifying the embedding function and distance metric
+    # Using cosine distance is common for text similarity tasks.
     collection = chroma_client.get_or_create_collection(
         name=COLLECTION_NAME,
         embedding_function=embedding_func,
-        metadata={"hnsw:space": "cosine"} # Use cosine distance
+        metadata={"hnsw:space": "cosine"} # Specify cosine distance metric
     )
+    st.success(f"✅ Chroma DB collection '{COLLECTION_NAME}' loaded/created at '{CHROMA_PATH}'.")
 except Exception as e:
-    st.error(f"Error initializing Chroma DB at '{CHROMA_PATH}': {e}")
-    st.info("If this is the first run, the directory will be created.")
-    # Attempt creation again more robustly if needed, or guide user.
+    st.error(f"❌ Error initializing Chroma DB at '{CHROMA_PATH}': {e}")
+    st.info("ℹ️ If this is the first run, the 'chroma_data' directory will be created.")
     st.stop()
 
 
 # --- Helper Functions ---
-def analyze_image_with_gemini(image_bytes):
-    """Sends image bytes to Gemini Vision and returns the text description."""
+
+def analyze_image_with_gemini(image_bytes: bytes) -> str:
+    """
+    Sends image bytes to the Gemini Vision model for analysis and returns
+    the generated text description.
+
+    Args:
+        image_bytes: The image data as bytes.
+
+    Returns:
+        A string containing the analysis text, or an error/blocked message.
+    """
     try:
         img = Image.open(io.BytesIO(image_bytes))
-        prompt = """Analyze this medical image (could be a pathology slide, diagram, or other medical visual).
-        Describe the key visual features relevant to a medical professional.
-        Identify potential:
-        - Diseases or conditions suggested
-        - Pathological findings (e.g., cellular morphology, tissue structure, staining patterns)
-        - Cell types visible
-        - Relevant biomarkers (if inferrable from staining or morphology)
-        - Anatomical context (if clear)
-
-        Be concise and focus on visually evident information.
-        """
+        # Define the prompt for the vision model
+        prompt = """Analyze this medical image (e.g., pathology slide, diagram, scan).
+Describe the key visual features relevant to a medical context.
+Identify potential:
+- Diseases or conditions indicated
+- Pathological findings (e.g., cellular morphology, tissue structure, staining patterns)
+- Visible cell types
+- Relevant biomarkers (if inferable from staining or morphology)
+- Anatomical context (if discernible)
+
+Be concise and focus primarily on visually evident information. Avoid definitive diagnoses.
+"""
+        # Generate content using the model
         response = gemini_model.generate_content([prompt, img])
-        # Handle potential blocked responses or errors
-        if not response.parts:
-             # Check if it was blocked
-             if response.prompt_feedback and response.prompt_feedback.block_reason:
-                 return f"Analysis blocked: {response.prompt_feedback.block_reason}"
-             else:
-                 # Some other issue, maybe no response text?
-                 return "Error: Gemini analysis failed or returned no content."
 
+        # Check for blocked content or empty response
+        if not response.parts:
+            if response.prompt_feedback and response.prompt_feedback.block_reason:
+                block_reason = response.prompt_feedback.block_reason
+                st.warning(f"⚠️ Analysis blocked by safety settings: {block_reason}")
+                return f"Analysis blocked due to safety settings: {block_reason}"
+            else:
+                st.error("❌ Gemini analysis returned no content. Response might be empty or invalid.")
+                return "Error: Gemini analysis failed or returned no content."
+
+        # Return the generated text
         return response.text
+
     except genai.types.BlockedPromptException as e:
-        st.error(f"Gemini request blocked: {e}")
-        return f"Analysis blocked due to safety settings: {e}"
+        st.error(f"❌ Gemini request blocked due to prompt content: {e}")
+        return f"Analysis blocked (prompt issue): {e}"
     except Exception as e:
-        st.error(f"Error during Gemini analysis: {e}")
+        st.error(f"❌ An error occurred during Gemini analysis: {e}")
         return f"Error analyzing image: {e}"
 
 
-def query_chroma(query_text, n_results=5):
-    """Queries the Chroma collection with the given text."""
+def query_chroma(query_text: str, n_results: int = 5) -> Optional[Dict[str, List[Any]]]:
+    """
+    Queries the Chroma DB collection with the given text.
+
+    Args:
+        query_text: The text to use for the similarity search.
+        n_results: The maximum number of results to return.
+
+    Returns:
+        A dictionary containing the query results ('documents', 'metadatas',
+        'distances'), or None if an error occurs.
+    """
     try:
         results = collection.query(
             query_texts=[query_text],
             n_results=n_results,
-            include=['documents', 'metadatas', 'distances'] # Include distances for relevance sorting
+            include=['documents', 'metadatas', 'distances'] # Specify fields to include
         )
         return results
     except Exception as e:
-        st.error(f"Error querying Chroma DB: {e}")
+        st.error(f"❌ Error querying Chroma DB: {e}")
         return None
 
 def add_dummy_data_to_chroma():
-    """Adds some example medical text snippets to Chroma."""
+    """
+    Adds predefined example medical text snippets and metadata to the Chroma collection.
+    Checks if documents with the same text already exist before adding.
+    """
+    st.info("Attempting to add dummy data to Chroma DB...")
+
     # --- IMPORTANT ---
-    # In a real scenario, this data would come from processing actual medical documents
-    # (papers, reports) using a tool like Unstructured (as in the original article)
-    # or manual curation to extract text and METADATA, including IMAGE_IDs.
-    # The embeddings generated here MUST match the query embedding function.
-    st.info("Adding dummy data to Chroma DB...")
+    # In a real application, this data ingestion process would involve:
+    # 1. Parsing actual medical documents (research papers, clinical notes, textbooks).
+    # 2. Extracting relevant text chunks (e.g., using tools like Unstructured).
+    # 3. Extracting or associating meaningful METADATA (source, patient ID (anonymized),
+    #    image IDs linked to text, extracted entities like diseases/genes).
+    # 4. Generating embeddings using the SAME embedding function used for querying.
     docs = [
         "Figure 1A shows adenocarcinoma of the lung, papillary subtype. Note the glandular structures and nuclear atypia. TTF-1 staining was positive.",
         "Pathology slide 34B demonstrates high-grade glioma (glioblastoma) with significant necrosis and microvascular proliferation. Ki-67 index was high.",
@@ -139,96 +196,127 @@ def add_dummy_data_to_chroma():
         "Slide CJD-Sample-02: Spongiform changes characteristic of prion disease are evident in the cerebral cortex. Gliosis is also noted."
     ]
     metadatas = [
-        {"source": "Example Paper 1", "entities": {"DISEASES": ["adenocarcinoma", "lung cancer"], "PATHOLOGY_FINDINGS": ["glandular structures", "nuclear atypia", "papillary subtype"], "BIOMARKERS": ["TTF-1"]}, "IMAGE_ID": "fig_1a_adeno_lung.png"},
-        {"source": "Path Report 789", "entities": {"DISEASES": ["high-grade glioma", "glioblastoma"], "PATHOLOGY_FINDINGS": ["necrosis", "microvascular proliferation"], "BIOMARKERS": ["Ki-67"]}, "IMAGE_ID": "slide_34b_gbm.tiff"},
-        {"source": "Textbook Chapter 5", "entities": {"GENES": ["EGFR"], "DRUGS": ["tyrosine kinase inhibitors"], "DISEASES": ["non-small cell lung cancer"]}, "IMAGE_ID": "diagram_egfr_pathway.svg"},
-        {"source": "Path Report 101", "entities": {"DISEASES": ["chronic gastritis", "Helicobacter pylori infection"], "PATHOLOGY_FINDINGS": ["intestinal metaplasia"]}, "IMAGE_ID": "micrograph_h_pylori_gastritis.jpg"},
-        {"source": "Case Study CJD", "entities": {"DISEASES": ["prion disease"], "PATHOLOGY_FINDINGS": ["Spongiform changes", "Gliosis"], "ANATOMICAL_LOCATIONS": ["cerebral cortex"]}, "IMAGE_ID": "slide_cjd_sample_02.jpg"}
+        {"source": "Example Paper 1", "topic": "Lung Cancer Pathology", "entities": {"DISEASES": ["adenocarcinoma", "lung cancer"], "PATHOLOGY_FINDINGS": ["glandular structures", "nuclear atypia", "papillary subtype"], "BIOMARKERS": ["TTF-1"]}, "IMAGE_ID": "fig_1a_adeno_lung.png"},
+        {"source": "Path Report 789", "topic": "Brain Tumor Pathology", "entities": {"DISEASES": ["high-grade glioma", "glioblastoma"], "PATHOLOGY_FINDINGS": ["necrosis", "microvascular proliferation"], "BIOMARKERS": ["Ki-67"]}, "IMAGE_ID": "slide_34b_gbm.tiff"},
+        {"source": "Textbook Chapter 5", "topic": "Molecular Oncology Pathways", "entities": {"GENES": ["EGFR"], "DRUGS": ["tyrosine kinase inhibitors"], "DISEASES": ["non-small cell lung cancer"]}, "IMAGE_ID": "diagram_egfr_pathway.svg"},
+        {"source": "Path Report 101", "topic": "Gastrointestinal Pathology", "entities": {"DISEASES": ["chronic gastritis", "Helicobacter pylori infection"], "PATHOLOGY_FINDINGS": ["intestinal metaplasia"]}, "IMAGE_ID": "micrograph_h_pylori_gastritis.jpg"},
+        {"source": "Case Study CJD", "topic": "Neuropathology", "entities": {"DISEASES": ["prion disease"], "PATHOLOGY_FINDINGS": ["Spongiform changes", "Gliosis"], "ANATOMICAL_LOCATIONS": ["cerebral cortex"]}, "IMAGE_ID": "slide_cjd_sample_02.jpg"}
     ]
-    ids = [f"doc_{int(time.time())}_{i}" for i in range(len(docs))] # Unique IDs
+    # Generate unique IDs using timestamp + index to minimize collision chance in demo
+    ids = [f"doc_{int(time.time())}_{i}" for i in range(len(docs))]
 
     try:
-        # Check if docs with these exact texts already exist to avoid duplicates on rerun
-        existing = collection.get(where={"$or": [{"document": doc} for doc in docs]})
-        if not existing or not existing['ids']: # Only add if none exist
-             collection.add(
-                 documents=docs,
-                 metadatas=metadatas,
-                 ids=ids
-             )
-             st.success(f"Added {len(docs)} dummy documents to Chroma collection '{COLLECTION_NAME}'.")
+        # Check if documents with these exact texts already exist to avoid duplicates
+        existing_docs = collection.get(where={"$or": [{"document": doc} for doc in docs]}, include=[]) # Don't need full data, just check existence
+        if not existing_docs or not existing_docs.get('ids'):
+            collection.add(
+                documents=docs,
+                metadatas=metadatas,
+                ids=ids
+            )
+            st.success(f"✅ Added {len(docs)} dummy documents to Chroma collection '{COLLECTION_NAME}'.")
         else:
-             st.warning("Dummy data seems to already exist in the collection.")
+            st.warning("⚠️ Dummy data (based on document text) seems to already exist in the collection. No new data added.")
 
     except Exception as e:
-        st.error(f"Error adding dummy data to Chroma: {e}")
+        st.error(f"❌ Error adding dummy data to Chroma: {e}")
 
 
 # --- Streamlit UI ---
-st.set_page_config(layout="wide")
+st.set_page_config(layout="wide", page_title="Medical Image Analysis & RAG")
 st.title("⚕️ Medical Image Analysis & RAG")
-st.markdown("Upload a medical image (pathology slide, diagram, etc.). Gemini Vision will analyze it, and Chroma DB will retrieve related information from a knowledge base.")
+st.markdown("""
+Upload a medical image (e.g., pathology slide, diagram).
+Google Gemini Vision will analyze it, and Chroma DB will retrieve related text snippets
+from a simulated knowledge base based on the analysis.
+""")
 
-# Sidebar for controls
+# Sidebar for Controls
 with st.sidebar:
-    st.header("Controls")
-    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png", "tiff", "webp"])
-    if st.button("Load Dummy KB Data"):
+    st.header("⚙️ Controls")
+    uploaded_file = st.file_uploader(
+        "Choose an image...",
+        type=["jpg", "jpeg", "png", "tiff", "webp"],
+        help="Upload a medical image file."
+    )
+
+    st.divider() # Visual separator
+
+    if st.button("➕ Load Dummy KB Data", help="Add example text data to the Chroma vector database."):
          add_dummy_data_to_chroma()
-    st.info("Note: Chroma data persists in the Space's storage but is lost if the Space is reset/deleted.")
 
+    st.divider()
 
-# Main area for display
-if uploaded_file is not None:
-    # Read image bytes
-    image_bytes = uploaded_file.getvalue()
+    st.info(f"""
+    ℹ️ **Note:**
+    - Chroma data is stored in the '{CHROMA_PATH}' folder within the app's environment.
+    - This data persists across runs but **will be lost** if the hosting environment (e.g., Streamlit Cloud, Hugging Face Space) is reset or its storage is cleared.
+    - Ensure the Google API Key is set in Streamlit Secrets.
+    """)
 
-    # Display the uploaded image
-    st.image(image_bytes, caption=f"Uploaded Image: {uploaded_file.name}", use_column_width=False, width=400)
 
-    st.markdown("---")
+# Main Display Area
+col1, col2 = st.columns(2) # Create two columns for layout
+
+with col1:
+    st.subheader("🖼️ Uploaded Image")
+    if uploaded_file is not None:
+        # Read image bytes from the uploaded file
+        image_bytes = uploaded_file.getvalue()
+        # Display the uploaded image
+        st.image(image_bytes, caption=f"Uploaded: {uploaded_file.name}", use_column_width=True)
+    else:
+        st.info("Upload an image using the sidebar to begin.")
+
+with col2:
     st.subheader("🔬 Gemini Vision Analysis")
+    if uploaded_file is not None:
+        # Analyze image with Gemini when an image is uploaded
+        with st.spinner("🧠 Analyzing image with Gemini Vision... This may take a moment."):
+            analysis_text = analyze_image_with_gemini(image_bytes)
+
+        # Display analysis or error message
+        if analysis_text.startswith("Error:") or analysis_text.startswith("Analysis blocked"):
+             # Errors/blocks are already logged via st.error/st.warning in the helper function
+             st.markdown(f"**Analysis Status:** {analysis_text}") # Show status message
+        else:
+            st.markdown(analysis_text)
+
+            st.markdown("---") # Separator before RAG results
+            st.subheader("📚 Related Information (RAG via Chroma DB)")
 
-    # Analyze image with Gemini
-    with st.spinner("Analyzing image with Gemini Vision..."):
-        analysis_text = analyze_image_with_gemini(image_bytes)
+            # Query Chroma DB using the Gemini analysis text
+            with st.spinner("🔍 Searching knowledge base..."):
+                chroma_results = query_chroma(analysis_text, n_results=3) # Fetch top 3 results
+
+            if chroma_results and chroma_results.get('documents') and chroma_results['documents'][0]:
+                num_results = len(chroma_results['documents'][0])
+                st.success(f"✅ Found {num_results} related entries in the knowledge base:")
+
+                for i in range(num_results):
+                    doc = chroma_results['documents'][0][i]
+                    meta = chroma_results['metadatas'][0][i]
+                    dist = chroma_results['distances'][0][i]
+
+                    expander_title = f"Result {i+1} (Similarity Score: {1-dist:.4f}) - Source: {meta.get('source', 'N/A')}"
+                    with st.expander(expander_title):
+                        st.markdown("**Retrieved Text:**")
+                        st.markdown(f"> {doc}") # Use blockquote for text
+                        st.markdown("**Metadata:**")
+                        st.json(meta) # Display metadata nicely formatted
+
+                        # Highlight if the retrieved text references another image/asset
+                        if meta.get("IMAGE_ID"):
+                            st.info(f"ℹ️ This text chunk is associated with visual asset: `{meta['IMAGE_ID']}`")
+                            # In a more complex app, you could add logic here to fetch/display this related image if available.
+
+            elif chroma_results is not None: # Query ran successfully but found nothing
+                 st.warning("⚠️ No relevant information found in the knowledge base matching the image analysis.")
+            # Else case (chroma_results is None) implies an error occurred, handled by st.error in query_chroma
 
-    if analysis_text.startswith("Error:") or analysis_text.startswith("Analysis blocked:"):
-         st.error(analysis_text)
     else:
-        st.markdown(analysis_text)
-
-        st.markdown("---")
-        st.subheader("📚 Related Information from Knowledge Base (Chroma DB)")
-
-        # Query Chroma DB using the Gemini analysis text
-        with st.spinner("Querying Chroma DB..."):
-            chroma_results = query_chroma(analysis_text)
-
-        if chroma_results and chroma_results.get('documents') and chroma_results['documents'][0]:
-            st.success(f"Found {len(chroma_results['documents'][0])} related entries:")
-            for i in range(len(chroma_results['documents'][0])):
-                doc = chroma_results['documents'][0][i]
-                meta = chroma_results['metadatas'][0][i]
-                dist = chroma_results['distances'][0][i]
-
-                with st.expander(f"Result {i+1} (Distance: {dist:.4f}) - Source: {meta.get('source', 'N/A')}"):
-                    st.markdown("**Text:**")
-                    st.markdown(doc)
-                    st.markdown("**Metadata:**")
-                    st.json(meta) # Display all metadata nicely
-
-                    # Highlight if it references another image
-                    if meta.get("IMAGE_ID"):
-                        st.info(f"ℹ️ This text describes another visual asset: `{meta['IMAGE_ID']}`")
-                        # In a real app, you might fetch/display this image if available
-        elif chroma_results is not None: # Query ran but found nothing
-             st.warning("No relevant information found in the knowledge base for this analysis.")
-        else: # Error occurred during query
-             st.error("Failed to retrieve results from Chroma DB.")
-
-else:
-    st.info("Upload an image using the sidebar to start the analysis.")
+        st.info("Analysis will appear here once an image is uploaded.")
+
 
 st.markdown("---")
-st.markdown("Powered by Google Gemini, Chroma DB, and Streamlit.")
+st.markdown("<div style='text-align: center;'>Powered by Google Gemini, Chroma DB, and Streamlit</div>", unsafe_allow_html=True)