Update icd9_ui.py

icd9_ui.py

@@ -225,7 +225,184 @@
 #     else:
 #         st.error("Please enter a medical summary.")
 
-import os
+# import os
+# import torch
+# import pandas as pd
+# import streamlit as st
+# from PIL import Image
+# from transformers import LongformerTokenizer, LongformerForSequenceClassification
+# from phi.agent import Agent
+# from phi.model.google import Gemini
+# from phi.tools.duckduckgo import DuckDuckGo
+
+# # Load the fine-tuned ICD-9 model and tokenizer
+# model_path = "./clinical_longformer"
+# tokenizer = LongformerTokenizer.from_pretrained(model_path)
+# model = LongformerForSequenceClassification.from_pretrained(model_path)
+# model.eval()  # Set the model to evaluation mode
+
+# # Load the ICD-9 descriptions from CSV into a dictionary
+# icd9_desc_df = pd.read_csv("D_ICD_DIAGNOSES.csv")  # Adjust the path to your CSV file
+# icd9_desc_df['ICD9_CODE'] = icd9_desc_df['ICD9_CODE'].astype(str)  # Ensure ICD9_CODE is string type for matching
+# icd9_descriptions = dict(zip(icd9_desc_df['ICD9_CODE'].str.replace('.', ''), icd9_desc_df['LONG_TITLE']))  # Remove decimals in ICD9 code for matching
+
+# # ICD-9 code columns used during training
+# icd9_columns = [
+#     '038.9', '244.9', '250.00', '272.0', '272.4', '276.1', '276.2', '285.1', '285.9',
+#     '287.5', '305.1', '311', '36.15', '37.22', '37.23', '38.91', '38.93', '39.61',
+#     '39.95', '401.9', '403.90', '410.71', '412', '414.01', '424.0', '427.31', '428.0',
+#     '486', '496', '507.0', '511.9', '518.81', '530.81', '584.9', '585.9', '599.0',
+#     '88.56', '88.72', '93.90', '96.04', '96.6', '96.71', '96.72', '99.04', '99.15',
+#     '995.92', 'V15.82', 'V45.81', 'V45.82', 'V58.61'
+# ]
+
+# # Function for making ICD-9 predictions
+# def predict_icd9(texts, tokenizer, model, threshold=0.5):
+    # inputs = tokenizer(
+    #     texts,
+    #     padding="max_length",
+    #     truncation=True,
+    #     max_length=512,
+    #     return_tensors="pt"
+    # )
+    
+    # with torch.no_grad():
+    #     outputs = model(
+    #         input_ids=inputs["input_ids"],
+    #         attention_mask=inputs["attention_mask"]
+    #     )
+    #     logits = outputs.logits
+    #     probabilities = torch.sigmoid(logits)
+    #     predictions = (probabilities > threshold).int()
+    
+    # predicted_icd9 = []
+    # for pred in predictions:
+    #     codes = [icd9_columns[i] for i, val in enumerate(pred) if val == 1]
+    #     predicted_icd9.append(codes)
+    
+    # predictions_with_desc = []
+    # for codes in predicted_icd9:
+    #     code_with_desc = [(code, icd9_descriptions.get(code.replace('.', ''), "Description not found")) for code in codes]
+    #     predictions_with_desc.append(code_with_desc)
+    
+    # return predictions_with_desc
+
+# Streamlit UI
+st.title("Medical Diagnosis Assistant")
+option = st.selectbox(
+    "Choose Diagnosis Method",
+    ("ICD-9 Code Prediction", "Medical Image Analysis")
+)
+
+# ICD-9 Code Prediction
+if option == "ICD-9 Code Prediction":
+    st.write("### Enter Medical Summary")
+    input_text = st.text_area("Medical Summary", placeholder="Enter clinical notes here...")
+
+    threshold = st.slider("Prediction Threshold", 0.0, 1.0, 0.5, 0.01)
+
+    if st.button("Predict ICD-9 Codes"):
+        if input_text.strip():
+            predictions = predict_icd9([input_text], tokenizer, model, threshold)
+            st.write("### Predicted ICD-9 Codes and Descriptions")
+            for code, description in predictions[0]:
+                st.write(f"- {code}: {description}")
+        else:
+            st.error("Please enter a medical summary.")
+
+# Medical Image Analysis
+# elif option == "Medical Image Analysis":
+#     if "GOOGLE_API_KEY" not in st.session_state:
+#         st.warning("Please enter your Google API Key in the sidebar to continue")
+#     else:
+#         medical_agent = Agent(
+#             model=Gemini(
+#                 api_key=st.session_state.GOOGLE_API_KEY,
+#                 id="gemini-2.0-flash-exp"
+#             ),
+#             tools=[DuckDuckGo()],
+#             markdown=True
+#         )
+
+#         query = """
+#         You are a highly skilled medical imaging expert with extensive knowledge in radiology and diagnostic imaging. Analyze the patient's medical image and structure your response as follows:
+
+#         ### 1. Image Type & Region
+#         - Specify imaging modality (X-ray/MRI/CT/Ultrasound/etc.)
+#         - Identify the patient's anatomical region and positioning
+#         - Comment on image quality and technical adequacy
+
+#         ### 2. Key Findings
+#         - List primary observations systematically
+#         - Note any abnormalities in the patient's imaging with precise descriptions
+#         - Include measurements and densities where relevant
+#         - Describe location, size, shape, and characteristics
+#         - Rate severity: Normal/Mild/Moderate/Severe
+
+#         ### 3. Diagnostic Assessment
+#         - Provide primary diagnosis with confidence level
+#         - List differential diagnoses in order of likelihood
+#         - Support each diagnosis with observed evidence from the patient's imaging
+        # - Note any critical or urgent findings
+
+        # ### 4. Patient-Friendly Explanation
+        # - Explain the findings in simple, clear language that the patient can understand
+        # - Avoid medical jargon or provide clear definitions
+        # - Include visual analogies if helpful
+        # - Address common patient concerns related to these findings
+
+        # ### 5. Research Context
+        # - Use the DuckDuckGo search tool to find recent medical literature about similar cases
+        # - Provide a list of relevant medical links
+        # - Include key references to support your analysis
+        # """
+
+        # upload_container = st.container()
+        # image_container = st.container()
+        # analysis_container = st.container()
+
+        # with upload_container:
+        #     uploaded_file = st.file_uploader(
+        #         "Upload Medical Image",
+        #         type=["jpg", "jpeg", "png", "dicom"],
+        #         help="Supported formats: JPG, JPEG, PNG, DICOM"
+        #     )
+
+        # if uploaded_file is not None:
+        #     with image_container:
+        #         col1, col2, col3 = st.columns([1, 2, 1])
+        #         with col2:
+        #             image = Image.open(uploaded_file)
+        #             width, height = image.size
+        #             aspect_ratio = width / height
+        #             new_width = 500
+        #             new_height = int(new_width / aspect_ratio)
+        #             resized_image = image.resize((new_width, new_height))
+                    
+        #             st.image(resized_image, caption="Uploaded Medical Image", use_container_width=True)
+                    
+        #             analyze_button = st.button("🔍 Analyze Image")
+
+        #     with analysis_container:
+        #         if analyze_button:
+        #             image_path = "temp_medical_image.png"
+        #             with open(image_path, "wb") as f:
+        #                 f.write(uploaded_file.getbuffer())
+                    
+        #             with st.spinner("🔄 Analyzing image... Please wait."):
+        #                 try:
+        #                     response = medical_agent.run(query, images=[image_path])
+        #                     st.markdown("### 📋 Analysis Results")
+        #                     st.markdown(response.content)
+        #                 except Exception as e:
+        #                     st.error(f"Analysis error: {e}")
+        #                 finally:
+        #                     if os.path.exists(image_path):
+        #                         os.remove(image_path)
+        # else:
+        #     st.info("👆 Please upload a medical image to begin analysis")
+
+        import os
 import torch
 import pandas as pd
 import streamlit as st
@@ -235,6 +412,16 @@ from phi.agent import Agent
 from phi.model.google import Gemini
 from phi.tools.duckduckgo import DuckDuckGo
 
+# Sidebar for Google API Key input
+st.sidebar.title("Settings")
+st.sidebar.write("Enter your Google API Key below for the Medical Image Analysis feature.")
+api_key = st.sidebar.text_input("Google API Key", type="password")
+
+if api_key:
+    st.session_state["GOOGLE_API_KEY"] = api_key
+else:
+    st.session_state.pop("GOOGLE_API_KEY", None)
+
 # Load the fine-tuned ICD-9 model and tokenizer
 model_path = "./clinical_longformer"
 tokenizer = LongformerTokenizer.from_pretrained(model_path)
@@ -317,7 +504,7 @@ elif option == "Medical Image Analysis":
     else:
         medical_agent = Agent(
             model=Gemini(
-                api_key=st.session_state.GOOGLE_API_KEY,
+                api_key=st.session_state["GOOGLE_API_KEY"],
                 id="gemini-2.0-flash-exp"
             ),
             tools=[DuckDuckGo()],
@@ -326,31 +513,26 @@ elif option == "Medical Image Analysis":
 
         query = """
         You are a highly skilled medical imaging expert with extensive knowledge in radiology and diagnostic imaging. Analyze the patient's medical image and structure your response as follows:
-
         ### 1. Image Type & Region
         - Specify imaging modality (X-ray/MRI/CT/Ultrasound/etc.)
         - Identify the patient's anatomical region and positioning
         - Comment on image quality and technical adequacy
-
         ### 2. Key Findings
         - List primary observations systematically
         - Note any abnormalities in the patient's imaging with precise descriptions
         - Include measurements and densities where relevant
         - Describe location, size, shape, and characteristics
         - Rate severity: Normal/Mild/Moderate/Severe
-
         ### 3. Diagnostic Assessment
         - Provide primary diagnosis with confidence level
         - List differential diagnoses in order of likelihood
         - Support each diagnosis with observed evidence from the patient's imaging
         - Note any critical or urgent findings
-
         ### 4. Patient-Friendly Explanation
         - Explain the findings in simple, clear language that the patient can understand
         - Avoid medical jargon or provide clear definitions
         - Include visual analogies if helpful
         - Address common patient concerns related to these findings
-
         ### 5. Research Context
         - Use the DuckDuckGo search tool to find recent medical literature about similar cases
         - Provide a list of relevant medical links
@@ -403,3 +585,4 @@ elif option == "Medical Image Analysis":
             st.info("👆 Please upload a medical image to begin analysis")
 
 
+