Update app.py

app.py

@@ -1,5 +1,4 @@
 import streamlit as st
-from streamlit.session_state import SessionState
 import pickle
 import pandas as pd
 import torch
@@ -12,6 +11,7 @@ from docx import Document
 from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
 from io import BytesIO
 import re
+from session_state import get
 
 openai.api_key = "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"
 
@@ -57,34 +57,19 @@ 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)
 
-# New function to handle feedback system
-def handle_feedback(session_state):
-    if session_state.caption_index < len(session_state.generated_captions) - 1:
-        session_state.caption_index += 1
-        new_caption = session_state.generated_captions[session_state.caption_index]
-        return new_caption
-    else:
-        st.warning("No more captions available to regenerate the report.")
-        return session_state.generated_captions[session_state.caption_index]
-
-# Initialize the session state
-if "caption_index" not in st.session_state:
-    st.session_state.caption_index = 0
-
-if "generated_captions" not in st.session_state:
-    st.session_state.generated_captions = []
+session_state = get(caption_index=0)
 
 def download_link(content, filename, link_text):
     b64 = base64.b64encode(content).decode()
     href = f'<a href="data:application/octet-stream;base64,{b64}" download="{filename}">{link_text}</a>'
     return href
 
-def show_predicted_caption(image, top_k=8):
+def show_predicted_caption_with_index(image, caption_index):
     matches = predict_caption(
         image, model, text_embeddings, testing_df["caption"]
-    )[:top_k]
-    cleaned_matches = [re.sub(r'\s\(ROCO_\d+\)', '', match) for match in matches]  # Add this line to clean the matches
-    return cleaned_matches  # Return the cleaned_matches instead of matches
+    )
+    cleaned_matches = [re.sub(r'\s\(ROCO_\d+\)', '', match) for match in matches]
+    return cleaned_matches[caption_index]
 
 def generate_radiology_report(prompt):
     response = openai.Completion.create(
@@ -109,49 +94,31 @@ def save_as_docx(text, filename):
         output.seek(0)
         return output.getvalue()
 
-st.title("RadiXGPT: An Evolution of machine doctors towards Radiology")
+st.title("RadiXGPT: An Evolution in Radiology Reporting")
+st.write("Upload a radiology image and generate a radiology report.")
 
-# Collect user's personal information
-st.subheader("Personal Information")
-first_name = st.text_input("First Name")
-last_name = st.text_input("Last Name")
-age = st.number_input("Age", min_value=0, max_value=120, value=25, step=1)
-gender = st.selectbox("Gender", ["Male", "Female", "Other"])
+uploaded_image = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
 
-st.write("Upload Scan to get Radiological Report:")
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
-if uploaded_file is not None:
-    image = Image.open(uploaded_file)
+if uploaded_image is not None:
+    image = Image.open(uploaded_image)
     st.image(image, caption="Uploaded Image", use_column_width=True)
-    st.write("")
 
     if st.button("Generate Caption"):
-        with st.spinner("Generating caption..."):
-            image_np = np.array(image)
-            captions = show_predicted_caption(image_np, top_k=2)
-            st.session_state.generated_captions = captions
-            caption = captions[st.session_state.caption_index]
-    
-            st.success(f"Caption: {caption}")
-
-            # Generate the radiology report
-            radiology_report = generate_radiology_report(f"Write Complete Radiology Report for this with clinical info, subjective, Assessment, Finding, Impressions, Conclusion and more in proper order : {caption}")
-
-            # Add personal information to the radiology report
-            radiology_report_with_personal_info = f"Patient Name: {first_name} {last_name}\nAge: {age}\nGender: {gender}\n\n{radiology_report}"
-
-            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)
-
-            # Feedback buttons
-            st.header("Feedback")
-            if st.button("Not Satisfied"):
-                new_caption = handle_feedback(st.session_state)
-                if new_caption:
-                    # Generate new radiology report
-                    radiology_report = generate_radiology_report(f"Write Complete Radiology Report for this with clinical info, subjective, Assessment, Finding, Impressions, Conclusion and more in proper order : {new_caption}")
-                    radiology_report_with_personal_info = f"Patient Name: {first_name} {last_name}\nAge: {age}\nGender: {gender}\n\n{radiology_report}"
-                    st.header("New 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 New Report as DOCX"), unsafe_allow_html=True)
+        if session_state.caption_index < 5:
+            session_state.caption_index += 1
+        else:
+            session_state.caption_index = 0
+
+    caption = show_predicted_caption_with_index(image, session_state.caption_index)
+    st.write("Caption: ", caption)
+
+    report_prompt = f"Generate a detailed radiology report based on this caption: {caption}"
+    report = generate_radiology_report(report_prompt)
+    st.write("Report: ", report)
+
+    # Download report as a .docx file
+    docx_data = save_as_docx(report, "Radiology_Report.docx")
+    st.markdown(
+        download_link(docx_data, "Radiology_Report.docx", "Download Report as .docx"),
+        unsafe_allow_html=True,
+    )