Update app.py

app.py

@@ -4,7 +4,7 @@ import pandas as pd
 import torch
 from PIL import Image
 import numpy as np
-from main import predict_caption, CLIPModel, get_text_embeddings
+from main import predict_caption, CLIPModel, get_text_embeddings, get_alternative_caption
 import openai
 import base64
 from docx import Document
@@ -56,12 +56,6 @@ 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 reset_session_state():
-    st.session_state.caption = None
-    st.session_state.report = None
-    st.session_state.feedback = None
-    st.session_state.regenerated_report = None
-
 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>'
@@ -71,8 +65,8 @@ def show_predicted_caption(image, top_k=8):
     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]
-    return cleaned_matches
+    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
 
 def generate_radiology_report(prompt):
     response = openai.Completion.create(
@@ -84,9 +78,11 @@ def generate_radiology_report(prompt):
         temperature=1,
     )
     report = response.choices[0].text.strip()
+    # Remove reference string from the report
     report = re.sub(r'\(ROCO_\d+\)', '', report).strip()
     return report
 
+
 def save_as_docx(text, filename):
     document = Document()
     document.add_paragraph(text)
@@ -95,9 +91,6 @@ def save_as_docx(text, filename):
         output.seek(0)
         return output.getvalue()
 
-# Call reset_session_state() at the beginning of the script
-reset_session_state()
-
 st.title("RadiXGPT: An Evolution of machine doctors towards Radiology")
 
 # Collect user's personal information
@@ -114,38 +107,34 @@ if uploaded_file is not None:
     st.image(image, caption="Uploaded Image", use_column_width=True)
     st.write("")
 
-    if st.button("Generate Caption") or st.session_state.caption:
-        if not st.session_state.caption:
-            with st.spinner("Generating caption..."):
-                image_np = np.array(image)
-                st.session_state.caption = show_predicted_caption(image_np)[0]
-    
-        st.success(f"Caption: {st.session_state.caption}")
-    
-        if not st.session_state.report:
-            st.session_state.report = generate_radiology_report(f"Write Complete Radiology Report for this with clinical info, subjective, Assessment, Finding, Impressions, Conclusion and more in proper order : {st.session_state.caption}")
-    
-        radiology_report_with_personal_info = f"Patient Name: {first_name} {last_name}\nAge: {age}\nGender: {gender}\n\n{st.session_state.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)
-    
+    if st.button("Generate Caption"):
+        with st.spinner("Generating caption..."):
+            image_np = np.array(image)
+            caption = show_predicted_caption(image_np)[0]
+
+            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_options = ["Satisfied", "Not Satisfied"]
+        selected_feedback = st.radio("Please provide feedback on the generated report:", feedback_options)
+
+        if selected_feedback == "Not Satisfied":
+            if st.button("Regenerate Report"):
+                with st.spinner("Regenerating report..."):
+                    alternative_caption = get_alternative_caption(image_np, model, text_embeddings, testing_df["caption"])
+                    regenerated_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 : {alternative_caption}")
+
+                    regenerated_radiology_report_with_personal_info = f"Patient Name: {first_name} {last_name}\nAge: {age}\nGender: {gender}\n\n{regenerated_radiology_report}"
 
-        if not 'feedback' in st.session_state:
-            st.session_state.feedback = feedback_options[0]
-        st.session_state.feedback = st.radio("Please provide feedback on the generated report:", feedback_options, key="feedback")
-
-        if st.session_state.feedback == "Not Satisfied":
-            if st.button("Regenerate Report") or st.session_state.regenerated_report:
-                if not st.session_state.regenerated_report:
-                    with st.spinner("Regenerating report..."):
-                        alternative_caption = get_alternative_caption(image_np, model, text_embeddings, testing_df["caption"])
-                        st.session_state.regenerated_report = generate_radiology_report(f"Write Complete Radiology Report for this with clinical info, subjective, Assessment, Finding, Impressions, Conclusion and more in proper order : {alternative_caption}")
-    
-                regenerated_radiology_report_with_personal_info = f"Patient Name: {first_name} {last_name}\nAge: {age}\nGender: {gender}\n\n{st.session_state.regenerated_report}"
-    
-                st.header("Regenerated Radiology Report")
-                st.write(regenerated_radiology_report_with_personal_info)
-                st.markdown(download_link(save_as_docx(regenerated_radiology_report_with_personal_info, "regenerated_radiology_report.docx"), "regenerated_radiology_report.docx", "Download Regenerated Report as DOCX"), unsafe_allow_html=True)
+                    st.header("Regenerated Radiology Report")
+                    st.write(regenerated_radiology_report_with_personal_info)
+                    st.markdown(download_link(save_as_docx(regenerated_radiology_report_with_personal_info, "regenerated_radiology_report.docx"), "regenerated_radiology_report.docx", "Download Regenerated Report as DOCX"), unsafe_allow_html=True)