Update icd9_ui.py

icd9_ui.py

@@ -1,66 +1,66 @@
-import streamlit as st
-import torch
-from transformers import LongformerTokenizer, LongformerForSequenceClassification
+# import streamlit as st
+# import torch
+# from transformers import LongformerTokenizer, LongformerForSequenceClassification
 
-# Load the fine-tuned 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 fine-tuned 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
 
-# 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'
-]
+# # 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 predictions
-def predict_icd9(texts, tokenizer, model, threshold=0.5):
-    inputs = tokenizer(
-        texts,
-        padding="max_length",
-        truncation=True,
-        max_length=512,
-        return_tensors="pt"
-    )
+# # Function for making 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()
+#     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)
+#     predicted_icd9 = []
+#     for pred in predictions:
+#         codes = [icd9_columns[i] for i, val in enumerate(pred) if val == 1]
+#         predicted_icd9.append(codes)
     
-    return predicted_icd9
+#     return predicted_icd9
 
-# Streamlit UI
-st.title("ICD-9 Code Prediction")
-st.sidebar.header("Model Options")
-model_option = st.sidebar.selectbox("Select Model", [ "ClinicalLongformer"])
-threshold = st.sidebar.slider("Prediction Threshold", 0.0, 1.0, 0.5, 0.01)
-
-st.write("### Enter Medical Summary")
-input_text = st.text_area("Medical Summary", placeholder="Enter clinical notes here...")
-
-if st.button("Predict"):
-    if input_text.strip():
-        predictions = predict_icd9([input_text], tokenizer, model, threshold)
-        st.write("### Predicted ICD-9 Codes")
-        for code in predictions[0]:
-            st.write(f"- {code}")
-    else:
-        st.error("Please enter a medical summary.")
+# # Streamlit UI
+# st.title("ICD-9 Code Prediction")
+# st.sidebar.header("Model Options")
+# model_option = st.sidebar.selectbox("Select Model", [ "ClinicalLongformer"])
+# threshold = st.sidebar.slider("Prediction Threshold", 0.0, 1.0, 0.5, 0.01)
+
+# st.write("### Enter Medical Summary")
+# input_text = st.text_area("Medical Summary", placeholder="Enter clinical notes here...")
+
+# if st.button("Predict"):
+#     if input_text.strip():
+#         predictions = predict_icd9([input_text], tokenizer, model, threshold)
+#         st.write("### Predicted ICD-9 Codes")
+#         for code in predictions[0]:
+#             st.write(f"- {code}")
+#     else:
+#         st.error("Please enter a medical summary.")
 
 # import torch
 # import pandas as pd
@@ -584,118 +584,118 @@ if st.button("Predict"):
 # #         else:
 # #             st.info("👆 Please upload a medical image to begin analysis")
 
-# 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
+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 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")
-# 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
+# Load the ICD-9 descriptions from CSV into a dictionary
+icd9_desc_df = pd.read_csv("D_ICD_DIAGNOSES.csv")
+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'
-# ]
+# 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()
+# 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)
+    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)
+    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
+    return predictions_with_desc
 
-# # Define the API key directly in the code
-# GOOGLE_API_KEY = "AIzaSyA24A6egT3L0NAKkkw9QHjfoizp7cJUTaA"
+# Define the API key directly in the code
+GOOGLE_API_KEY = "AIzaSyA24A6egT3L0NAKkkw9QHjfoizp7cJUTaA"
 
-# # 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":
-#     medical_agent = Agent(
-#         model=Gemini(
-#             api_key=GOOGLE_API_KEY,
-#             id="gemini-2.0-flash-exp"
-#         ),
-#         tools=[DuckDuckGo()],
-#         markdown=True
-#     )
+# 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":
+    medical_agent = Agent(
+        model=Gemini(
+            api_key=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
+    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
@@ -710,49 +710,49 @@ if st.button("Predict"):
 #     - Include key references to support your analysis
 #     """
 
-#     upload_container = st.container()
-#     image_container = st.container()
-#     analysis_container = st.container()
+    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"
-#         )
+    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))
+    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)
+                st.image(resized_image, caption="Uploaded Medical Image", use_container_width=True)
                 
-#                 analyze_button = st.button("🔍 Analyze Image")
+                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 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")
+                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")