app.py

import streamlit as st
from streamlit.session_state import SessionState
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
import re

openai.api_key = "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"

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)

# 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>'
    return href

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]  # 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(
        engine="text-davinci-003",
        prompt=prompt,
        max_tokens=800,
        n=1,
        stop=None,
        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)
    with BytesIO() as output:
        document.save(output)
        output.seek(0)
        return output.getvalue()

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")
age = st.number_input("Age", min_value=0, max_value=120, value=25, step=1)
gender = st.selectbox("Gender", ["Male", "Female", "Other"])

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)
    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)