merge code

src/streamlit_app.py

@@ -107,6 +107,7 @@ TEXT_COLUMN = "complaints"  # or "general_complaint", depending on your needs
 st.title("🩺 Multimodal Medical Chatbot")
 
 query = st.text_input("Enter your medical question or symptom description:")
+uploaded_file = st.file_uploader("Upload an image to find similar medical cases:", type=["png", "jpg", "jpeg"])
 
 @track
 def get_chat_completion_openai(client, prompt: str):
@@ -127,6 +128,21 @@ def get_similar_prompt(query):
     idx = top_result.indices[0].item()
     return data[idx]
 
+# Cache dataset image embeddings (takes time, so cached)
+@st.cache_data(show_spinner=True)
+def embed_dataset_images(dataset):
+    features = []
+    for item in dataset:
+        # Load image from URL/path or raw bytes - adapt this if needed
+        img = Image.open(item["image"]).convert("RGB")
+        inputs = clip_processor(images=img, return_tensors="pt")
+        with torch.no_grad():
+            feat = clip_model.get_image_features(**inputs)
+        feat /= feat.norm(p=2, dim=-1, keepdim=True)
+        features.append(feat.cpu())
+    return torch.cat(features, dim=0)
+
+dataset_image_features = embed_dataset_images(data)
 
 if query:
     with st.spinner("Searching medical cases..."):
@@ -152,4 +168,28 @@ if query:
         else:
             st.warning("OpenAI API key not found. Please set OPENAI_API_KEY as a secret environment variable.")
 
+if uploaded_file is not None:
+    query_image = Image.open(uploaded_file).convert("RGB")
+    st.image(query_image, caption="Your uploaded image", use_container_width=True)
+
+    # Embed uploaded image
+    inputs = clip_processor(images=query_image, return_tensors="pt")
+    with torch.no_grad():
+        query_feat = clip_model.get_image_features(**inputs)
+    query_feat /= query_feat.norm(p=2, dim=-1, keepdim=True)
+
+    # Compute cosine similarity
+    similarities = (dataset_image_features @ query_feat.T).squeeze(1)  # [num_dataset_images]
+
+    top_k = 3
+    top_results = torch.topk(similarities, k=top_k)
+
+    st.write(f"Top {top_k} similar medical cases:")
+
+    for rank, idx in enumerate(top_results.indices):
+        score = top_results.values[rank].item()
+        similar_img = Image.open(data[int(idx)]["image"]).convert("RGB")
+        st.image(similar_img, caption=f"Similarity: {score:.3f}", use_container_width=True)
+        st.markdown(f"**Case description:** {data[int(idx)]['complaints']}")
+
 st.caption("This chatbot is for educational purposes only and does not provide medical advice.")