import torch
from PIL import Image
import streamlit as st
import numpy as np
import pandas as pd
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 = "sk-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"
st.markdown(
"""
""",
unsafe_allow_html=True,
)
device = torch.device("cpu")
testing_df = pd.read_csv("testing_df.csv")
model = CLIPModel() # Create an instance of CLIPModel
# Load the model
state_dict = torch.load("weights.pt", map_location=torch.device('cpu'))
print("Loaded State Dict Keys:", state_dict.keys())
# Create an instance of CLIPModel
model = CLIPModel().to(device)
print("Model Keys:", model.state_dict().keys())
# Load the state_dict into the model
model.load_state_dict(state_dict, strict=False) # Set strict=False to ignore unexpected keys
text_embeddings = torch.load('saved_text_embeddings.pt', map_location=device)
def download_link(content, filename, link_text):
b64 = base64.b64encode(content).decode()
href = f'{link_text}'
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)
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}"
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)