Update app.py

app.py

@@ -1,4 +1,5 @@
 import streamlit as st
+from streamlit.session_state import SessionState
 import pickle
 import pandas as pd
 import torch
@@ -56,6 +57,23 @@ 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 = []
+
 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>'
@@ -110,8 +128,10 @@ 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)[0]
-
+            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
@@ -125,10 +145,13 @@ if uploaded_file is not None:
             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("Thanks for your feedback!")
-            feedback_options = ["Better", "Satisfied", "Worse"]
-            selected_feedback = st.radio("Please provide feedback on the generated report:", feedback_options)
-
-            if st.button("Submit Feedback"):
-                st.success("Thanks for providing feedback!")
-                # Implement your feedback handling logic here based on the `selected_feedback` value
+            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)