Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,108 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+import os, torch, transformers
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from PIL import Image
+from torchvision import transforms
+from io import BytesIO
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+@st.cache_resource
+def load_caption_model():   
+    # load medicap
+    ckpt_name = 'aehrc/medicap'
+    medicap = transformers.AutoModel.from_pretrained(ckpt_name, trust_remote_code=True)
+    medicap = medicap.to(device)
+    medicap.eval()
+
+    # transform image 
+    image_processor = transformers.AutoFeatureExtractor.from_pretrained(ckpt_name)
+    medicap_transforms = transforms.Compose(
+        [
+            transforms.Resize(size=image_processor.size['shortest_edge']),
+            transforms.CenterCrop(size=[
+                image_processor.size['shortest_edge'],
+                image_processor.size['shortest_edge'],
+            ]
+            ),
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=image_processor.image_mean,
+                std=image_processor.image_std,
+            ),
+        ]
+    )
+
+    # tokenizer
+    medicap_tokenizer = transformers.GPT2Tokenizer.from_pretrained(ckpt_name)
+
+    return medicap, medicap_transforms, medicap_tokenizer
+
+def generate_image_caption(image, model, transformer, tokenizer):
+    image = transformer(image)
+    image = image.unsqueeze(0)
+    outputs = model.generate(
+        pixel_values=image.to(device),
+        bos_token_id=tokenizer.bos_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.pad_token_id,
+        return_dict_in_generate=True,
+        use_cache=True,
+        max_length=256,
+        num_beams=4,
+        output_attentions=False
+    )
+    return tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
+
+@st.cache_resource
+def load_qa_model():
+    model_name = "microsoft/BioGPT-Large-PubMedQA"
+    biogpt_tokenizer = AutoTokenizer.from_pretrained(model_name)
+    biogpt = AutoModelForCausalLM.from_pretrained(model_name)
+    biogpt = biogpt.to(device)
+    biogpt.eval()
+
+    return biogpt, biogpt_tokenizer
+
+def generate_answer(description, question, model, tokenizer):
+    prompt = f"question: {question} context: {description}"
+    new_input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
+    input_ids = new_input_ids
+
+    generated_output = model.generate(
+        input_ids,
+        max_new_tokens=100,  # Max new tokens for the bot's response
+    )
+
+    response = tokenizer.decode(generated_output[0], skip_special_tokens=True)
+
+    return response
+
+st.set_page_config(page_title="Image Caption + QA", layout="centered")
+st.title("🖼️ Caption-Based Question Answering")
+
+uploaded_file = st.file_uploader("Choose Image", type = ["jpg", "jpeg", "png"])
+if uploaded_file is not None:
+    img = Image.open(uploaded_file)
+    st.image(img)
+
+    # image description
+    medicap, medicap_transforms, medicap_tokenizer = load_caption_model()
+    caption = generate_image_caption(img, medicap, medicap_transforms, medicap_tokenizer)
+
+    st.subheader("📝 Generated Description")
+    st.info(caption)
+
+    # vqa
+    st.markdown("---")
+
+    st.subheader("❓ Ask a Question About the Image")
+    question = st.text_input("Type your question")
+
+    if question:
+        biogpt, biogpt_tokenizer = load_qa_model()
+        response = generate_answer(caption, question, biogpt, biogpt_tokenizer)
+        st.success(f"{response}")
+
+else:
+    st.info("Please upload an image file.")