Update app.py

app.py

@@ -1,85 +1,29 @@
 import streamlit as st
-import pickle
-import pandas as pd
-import torch
 from PIL import Image
 import numpy as np
-from main import predict_caption, CLIPModel, get_text_embeddings
-import openai
-import base64
-from docx import Document
-from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
-from io import BytesIO
-
-# Set up OpenAI API
-openai.api_key = "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"
-
-# Custom CSS for the page
-st.markdown(
-    """
-<style>
-    body {
-        background-color: transparent;
-    }
-    .container {
-        display: flex;
-        justify-content: center;
-        align-items: center;
-        background-color: rgba(255, 255, 255, 0.7);
-        border-radius: 15px;
-        padding: 20px;
-    }
-    .stApp {
-        background-color: transparent;
-    }
-    .stText, .stMarkdown, .stTextInput>label, .stButton>button>span {
-        color: #1c1c1c !important; /* Set the dark text color for text elements */
-    }
-    .stButton>button>span {
-        color: initial !important; /* Reset the text color for the 'Generate Caption' button */
-    }
-    .stMarkdown h1, .stMarkdown h2 {
-        color: #ff6b81 !important; /* Set the text color of h1 and h2 elements to soft red-pink */
-        font-weight: bold; /* Set the font weight to bold */
-        border: 2px solid #ff6b81; /* Add a bold border around the headers */
-        padding: 10px; /* Add padding to the headers */
-        border-radius: 5px; /* Add border-radius to the headers */
-    }
-</style>
-""",
-    unsafe_allow_html=True,
-)
-
-device = torch.device("cpu")
-
-testing_df = pd.read_csv("testing_df.csv")
-model = CLIPModel().to(device)
-model.load_state_dict(torch.load("weights.pt", map_location=torch.device('cpu')))
-text_embeddings = torch.load('saved_text_embeddings.pt', map_location=device)
-
-def show_predicted_caption(image, second_attempt=False):
-    matches = predict_caption(
-        image, model, text_embeddings, testing_df["caption"], second_attempt
-    )[0 if not second_attempt else 1]
-    return matches
-
-def generate_radiology_report(prompt):
-    response = openai.Completion.create(
-        engine="text-davinci-003",
-        prompt=prompt,
-        max_tokens=800,
-        n=1,
-        stop=None,
-        temperature=1,
-    )
-    return response.choices[0].text.strip()
-
-st.title("RadiXGPT: An Evolution of machine doctors towards Radiology")
-
-# Collect user's personal information
-st.subheader("Personal Information")
-first_name = st.text_input("First Name")
-last_name = st.text_input("Last Name")
+import main
+import pandas as pd
+import cv2
+
+def get_image_np(uploaded_file):
+    image = Image.open(uploaded_file)
+    return np.array(image)
+
+def show_predicted_caption(image_np, top_k=1):
+    image = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
+    captions = predict_caption(image, model, text_embeddings, valid_df['caption'].values, n=top_k)
+    return captions
+
+# Load the model and text embeddings
+valid_df = pd.read_csv('testing_df.csv')
+model, text_embeddings = main.get_text_embeddings(valid_df)
+
+# App code
+st.title("Medical Radiology Report Generator")
+
+st.header("Personal Information")
+first_name = st.text_input("First Name", "John")
+last_name = st.text_input("Last Name", "Doe")
 age = st.number_input("Age", min_value=0, max_value=120, value=25, step=1)
 gender = st.selectbox("Gender", ["Male", "Female", "Other"])
 
@@ -93,7 +37,7 @@ if uploaded_file is not None:
     if st.button("Generate Caption"):
         with st.spinner("Generating caption..."):
             image_np = np.array(image)
-            caption = show_predicted_caption(image_np)
+            caption = show_predicted_caption(image_np)[0]
 
             st.success(f"Caption: {caption}")
 
@@ -109,20 +53,24 @@ if uploaded_file is not None:
                 st.write(radiology_report_with_personal_info)
                 st.markdown(download_link(save_as_docx(radiology_report_with_personal_info, "radiology_report.docx"), "radiology_report.docx", "Download Report as DOCX"), unsafe_allow_html=True)
 
-                st.header("Self Adaptive Feedback System")
+            # Advanced Feedback System
+            st.header("Advanced Feedback System")
+            feedback_options = ["Better", "Satisfied", "Worse"]
+            feedback = st.radio("Rate the generated report:", feedback_options)
 
-                feedback_options = st.radio("Is the radiology report satisfactory?", ["Better", "Satisfied", "Worse"])
+            top_k = 1
+            while feedback == "Worse":
+                top_k += 1
+                with st.spinner("Regenerating report..."):
+                    new_caption = show_predicted_caption(image_np, top_k=top_k)[-1]
+                    radiology_report = generate_radiology_report(f"Write Complete Radiology Report for this: {new_caption}")
 
-                if feedback_options == "Worse":
-                    caption = show_predicted_caption(image_np, second_attempt=True)
-                    radiology_report = generate_radiology_report(f"Write Complete Radiology Report for this: {caption}")
                     radiology_report_with_personal_info = f"Patient Name: {first_name} {last_name}\nAge: {age}\nGender: {gender}\n\n{radiology_report}"
+
+                with container:
+                    container.empty()
+                    st.header("Radiology Report")
                     st.write(radiology_report_with_personal_info)
                     st.markdown(download_link(save_as_docx(radiology_report_with_personal_info, "radiology_report.docx"), "radiology_report.docx", "Download Report as DOCX"), unsafe_allow_html=True)
-                elif feedback_options == "Satisfied":
-                    st.success("Thank you for your feedback! Your input helps us improve the RadiXGPT system.")
-                elif feedback_options == "Better":
-                    st.success("Thank you for your feedback! Your input helps us improve the RadiXGPT system.")
-                # Add your feedback processing logic here
-                # This could include updating the model with the new information,
-                # logging the feedback, or triggering an alert for manual review
+
+                feedback = st.radio("Rate the generated report:", feedback_options)