Remove binary files from repository

.gitignore

@@ -1,2 +1,3 @@
 venv
-.env
+.env
+test_model.py

app.py

@@ -1,7 +1,10 @@
 import streamlit as st
 from transformers import pipeline
 from PIL import Image
+import numpy as np
+import cv2
 from rag_model import *
+from yolo_model import *
 
 
 @st.cache_resource
@@ -9,64 +12,209 @@ def load_image_model():
     return pipeline("image-classification", model="Heem2/wound-image-classification")
 
 pipeline = load_image_model()
+yolo_model = load_yolo_model()
 
-st.title("FirstAid-AI")
+# Add custom CSS
+css = """
+<style>
+body {
+    font-family: 'Arial', sans-serif;
+    background-color: #f5f5f5;
+}
+.main {
+    background-color: #ffffff;
+    padding: 20px;
+    border-radius: 10px;
+    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
+}
+.stButton button {
+    background-color: #4CAF50;
+    color: white;
+    border: none;
+    padding: 10px 20px;
+    text-align: center;
+    text-decoration: none;
+    display: inline-block;
+    font-size: 16px;
+    margin: 4px 2px;
+    cursor: pointer;
+    border-radius: 5px;
+}
+.stButton button:hover {
+    background-color: #45a049;
+}
+.stApp > header {
+    background-color: transparent;
+}
+.stApp {
+    margin: auto;
+    background-color: #D9AFD9;
+    background-image: linear-gradient(0deg, #D9AFD9 0%, #97D9E1 100%);
+}
+[data-testid='stFileUploader'] {
+    width: max-content;
+}
+[data-testid='stFileUploader'] section {
+    padding: 0;
+    float: left;
+}
+[data-testid='stFileUploader'] section > input + div {
+    display: none;
+}
+[data-testid='stFileUploader'] section + div {
+    float: right;
+    padding-top: 0;
+}
+</style>
+"""
+
+st.markdown(css, unsafe_allow_html=True)
+
+st.title("**FirstAid-AI**")
+
+# Add a description at the top
+st.markdown("""
+### Welcome to FirstAid-AI
+This application provides medical advice based on images of wounds and medical equipment. 
+Upload an image of your wound or medical equipment, and the AI will classify the image and provide relevant advice.
+""")
+
+st.markdown("## How to Use FirstAid-AI")
+st.markdown("### 1. Upload an image of a wound and a piece of equipment (if applicable)")
+st.image("images/example3.png", use_container_width=True)
+st.caption("The AI model will detect the wound or equipment in the image and provide confidence levels. The AI assistant will then provide treatment or usage advice.")
+st.markdown("### 2. Ask follow-up questions and continue the conversation with the AI assistant!")
 
 # Initialize chat history
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# Upload an image of a wound
-file = st.file_uploader("Upload an image of your wound")
+# Dropdown to select the type of images to provide
+option = st.selectbox(
+    "Select the type of images you want to provide:",
+    ("Provide just wound image", "Provide both wound and equipment")
+)
+
+# Upload images based on the selected option
+file_wound = None
+file_equipment = None
+
+if option == "Provide just wound image":
+    file_wound = st.file_uploader("Upload an image of your wound")
+elif option == "Provide both wound and equipment":
+    file_wound = st.file_uploader("Upload an image of your wound")
+    file_equipment = st.file_uploader("Upload an image of your equipment")
 
 # Reset chat history if no file is uploaded
-if file is None:
+if file_wound is None and file_equipment is None:
     st.session_state.messages = []
 
-if file is not None:
-    # Display the image and predictions
+if file_wound is not None and option == "Provide just wound image":
+    # Display the wound image and predictions
     col1, col2 = st.columns(2)
-    image = Image.open(file)
+    image = Image.open(file_wound)
     col1.image(image, use_container_width=True)
 
-    # Classify the image
+    # Classify the wound image
     predictions = pipeline(image)
     detected_wound = predictions[0]['label']
     col2.header("Detected Wound")
     for p in predictions:
         col2.subheader(f"{p['label']}: {round(p['score'] * 100, 1)}%")
 
-    # Initial advice for wound advice
+    # Initial advice for wound
     if not st.session_state.messages:
         initial_query = f"Provide treatment advice for a {detected_wound} wound"
         initial_response = rag_chain.invoke(initial_query)
         st.session_state.messages.append({"role": "assistant", "content": initial_response})
 
-# Display chat messages from history
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-
-# Accept user input if an image is uploaded
-if file is not None and (prompt := st.chat_input("Ask a follow-up question or continue the conversation:")):
-    # Display user message in chat
-    with st.chat_message("user"):
-        st.markdown(prompt)
-    # Add user message to chat history
-    st.session_state.messages.append({"role": "user", "content": prompt})
-
-    # Prepare the conversation history for rag_chain
-    conversation_history = "\n".join(
-        f"{message['role']}: {message['content']}" for message in st.session_state.messages
-    )
-
-    # Generate response from rag_chain
-    query = f"Context:\n{conversation_history}\n\nAssistant, respond to the user's latest query: {prompt}"
-    response = rag_chain.invoke(query)
-
-    # Display assistant response in chat message container
-    with st.chat_message("assistant"):
-        st.markdown(response)
-        
-    # Add assistant response to chat history
-    st.session_state.messages.append({"role": "assistant", "content": response})
+    # Display chat messages from history
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+
+    # Accept user input if an image is uploaded
+    if (file_wound is not None or file_equipment is not None) and (prompt := st.chat_input("Ask a follow-up question or continue the conversation:")):
+        # Display user message in chat
+        with st.chat_message("user"):
+            st.markdown(prompt)
+        # Add user message to chat history
+        st.session_state.messages.append({"role": "user", "content": prompt})
+
+        # Prepare the conversation history for rag_chain
+        conversation_history = "\n".join(
+            f"{message['role']}: {message['content']}" for message in st.session_state.messages
+        )
+
+        # Generate response from rag_chain
+        query = f"Context:\n{conversation_history}\n\nAssistant, respond to the user's latest query: {prompt}"
+        response = rag_chain.invoke(query)
+
+        # Display assistant response in chat message container
+        with st.chat_message("assistant"):
+            st.markdown(response)
+
+        # Add assistant response to chat history
+        st.session_state.messages.append({"role": "assistant", "content": response})
+
+if file_wound is not None and file_equipment is not None and option == "Provide both wound and equipment":
+    # Display the wound image and predictions
+    col1, col2 = st.columns(2)
+    image = Image.open(file_wound)
+    col1.image(image, use_container_width=True)
+
+    # Classify the wound image
+    predictions = pipeline(image)
+    detected_wound = predictions[0]['label']
+    col2.header("Detected Wound")
+    for p in predictions:
+        col2.subheader(f"{p['label']}: {round(p['score'] * 100, 1)}%")
+    
+    # Display the equipment image and predictions
+    col3, col4 = st.columns(2)
+    image = Image.open(file_equipment)
+    col3.image(image, use_container_width=True)
+
+    # Convert the image to a format supported by YOLO
+    image_np = np.array(image)
+    image_cv = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
+
+    # Classify the equipment image using YOLO model
+    detected_equipment = get_detected_objects(yolo_model, image_cv)
+    col4.header("Detected Equipment")
+    col4.subheader(detected_equipment)
+
+    # Initial advice for equipment
+    if not st.session_state.messages:
+        initial_query = f"Provide usage advice for {detected_equipment} when treating a {detected_wound} wound"
+        initial_response = rag_chain.invoke(initial_query)
+        st.session_state.messages.append({"role": "assistant", "content": initial_response})
+
+    # Display chat messages from history
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+
+    # Accept user input if an image is uploaded
+    if (file_wound is not None or file_equipment is not None) and (prompt := st.chat_input("Ask a follow-up question or continue the conversation:")):
+        # Display user message in chat
+        with st.chat_message("user"):
+            st.markdown(prompt)
+        # Add user message to chat history
+        st.session_state.messages.append({"role": "user", "content": prompt})
+
+        # Prepare the conversation history for rag_chain
+        conversation_history = "\n".join(
+            f"{message['role']}: {message['content']}" for message in st.session_state.messages
+        )
+
+        # Generate response from rag_chain
+        query = f"Context:\n{conversation_history}\n\nAssistant, respond to the user's latest query: {prompt}"
+        response = rag_chain.invoke(query)
+
+        # Display assistant response in chat message container
+        with st.chat_message("assistant"):
+            st.markdown(response)
+
+        # Add assistant response to chat history
+        st.session_state.messages.append({"role": "assistant", "content": response})

models/best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:54c49363a8d6f9a503178ab1f0594b4ece9ea35f30cc8c2721d4564c63bbb48e
+size 22528035

rag_model.py

@@ -1,13 +1,11 @@
-# Import libraries.
-
-# File loading and environment variables.
+# File loading and environment variables
 import os
 from dotenv import load_dotenv
 
-# HuggingFace LLM.
-from huggingface_hub import InferenceClient
+# Gemini Library
+import google.generativeai as genai
 
-# Langchain.
+# Langchain
 from langchain.document_loaders import TextLoader
 from langchain.prompts import PromptTemplate
 from langchain.schema.runnable import RunnablePassthrough, RunnableLambda
@@ -15,34 +13,39 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import MongoDBAtlasVectorSearch
 from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings
 
-# MongoDB.
+# MongoDB
 from pymongo import MongoClient
 
-# Function type hints.
+# Function type hints
 from typing import Dict, Any
 
 # Streamlit
 import streamlit as st
 
-# Load the environment variables from the .env file
+# Load environment variables
 load_dotenv()
 
-# Load the environment variables
+# Retrieve environment variables
 MONGO_URI = os.getenv("MONGO_URI")
 HF_TOKEN = os.getenv("HF_TOKEN")
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+
+# Configure Gemini
+genai.configure(api_key=GEMINI_API_KEY)
+model = genai.GenerativeModel("gemini-1.5-flash")
 
-# Setup Vector Store and MongoDB Atlas connection.
+# Setup Vector Store and MongoDB Atlas connection
 
-# Connect to MongoDB Atlas cluster using the connection string.
+# Connect to MongoDB Atlas cluster using the connection string
 DB_NAME = "ericguan04"
 COLLECTION_NAME = "first_aid_intents"
 vector_search_index = "vector_index"
 
 @st.cache_resource
 def get_mongodb_collection():
-    # Connect to MongoDB Atlas cluster using the connection string.
+    # Connect to MongoDB Atlas cluster using the connection string
     cluster = MongoClient(MONGO_URI)
-    # Connect to the specific collection in the database.
+    # Connect to the specific collection in the database
     return cluster[DB_NAME][COLLECTION_NAME]
 
 MONGODB_COLLECTION = get_mongodb_collection()
@@ -62,25 +65,23 @@ vector_search = MongoDBAtlasVectorSearch.from_connection_string(
     index_name=vector_search_index,
 )
 
-# k to search for only the X most relevant documents.
+# k to search for only the X most relevant documents
 k = 10
 
-# score_threshold to use only documents with a relevance score above 0.80.
+# score_threshold to use only documents with a relevance score above 0.80
 score_threshold = 0.80
 
 # Build your retriever
 retriever_1 = vector_search.as_retriever(
-   search_type = "similarity", # similarity, mmr, similarity_score_threshold. https://api.python.langchain.com/en/latest/vectorstores/langchain_core.vectorstores.VectorStore.html#langchain_core.vectorstores.VectorStore.as_retriever
-   search_kwargs = {"k": k, "score_threshold": score_threshold}
+   search_type="similarity",  # similarity, mmr, similarity_score_threshold. https://api.python.langchain.com/en/latest/vectorstores/langchain_core.vectorstores.VectorStore.html#langchain_core.vectorstores.VectorStore.as_retriever
+   search_kwargs={"k": k, "score_threshold": score_threshold},
 )
 
-
-# Initialize Hugging Face client
-hf_client = InferenceClient(api_key=HF_TOKEN)
-
 # Define the prompt template
 prompt = PromptTemplate.from_template(
-    """Use the following pieces of context to answer the question at the end.
+    """You are playing the role of a medical assistant. A patient has come to you with a minor medical issue.
+    Use the following pieces of context to answer the question at the end.
+    To be more natural, do not mention you are referring to the context.
 
     START OF CONTEXT:
     {context}
@@ -92,41 +93,34 @@ prompt = PromptTemplate.from_template(
 
     If you do not know the answer, just say that you do not know.
     NEVER assume things.
+    If the question is not relevant to the context, just say that it is not relevant.
     """
 )
 
-
-# Formatting the retrieved documents before inserting them in the system prompt template.
+# Formatting the retrieved documents before inserting them in the system prompt template
 def format_docs(docs):
     return "\n\n".join(doc.page_content for doc in docs)
 
 @st.cache_resource
 def generate_response(input_dict: Dict[str, Any]) -> str:
+    """
+    Generate a response using the Gemini model.
+    
+    Parameters:
+        input_dict (Dict[str, Any]): Dictionary with formatted context and question.
+        
+    Returns:
+        str: Generated response from the Gemini model.
+    """
     formatted_prompt = prompt.format(**input_dict)
-    # print(formatted_prompt)
-
-    ## THIS IS YOUR LLM
-    response = hf_client.chat.completions.create(
-        model="Qwen/Qwen2.5-1.5B-Instruct",
-        messages=[{
-            "role": "system",
-            "content": formatted_prompt
-        },{
-            "role": "user",
-            "content": input_dict["question"]
-        }],
-        max_tokens=1000,
-        temperature=0.2,
-    )
-
-    return response.choices[0].message.content
+    response = model.generate_content(formatted_prompt)
+    return response.text  # Adjust based on actual response structure
 
-
-# Build the chain with retriever_1.
+# Build the chain with retriever_1
 rag_chain = (
     {
         "context": retriever_1 | RunnableLambda(format_docs),
-        "question": RunnablePassthrough()
+        "question": RunnablePassthrough(),
     }
     | RunnableLambda(generate_response)
-)
+)

requirements.txt

@@ -10,15 +10,28 @@ cachetools==5.5.0
 certifi==2024.8.30
 charset-normalizer==3.4.0
 click==8.1.7
+contourpy==1.3.1
+cycler==0.12.1
 dataclasses-json==0.6.7
 dnspython==2.7.0
 filelock==3.16.1
+fonttools==4.55.2
 frozenlist==1.5.0
 fsspec==2024.10.0
 gitdb==4.0.11
 GitPython==3.1.43
+google-ai-generativelanguage==0.6.10
+google-api-core==2.23.0
+google-api-python-client==2.154.0
+google-auth==2.36.0
+google-auth-httplib2==0.2.0
+google-generativeai==0.8.3
+googleapis-common-protos==1.66.0
+grpcio==1.68.1
+grpcio-status==1.68.1
 h11==0.14.0
 httpcore==1.0.7
+httplib2==0.22.0
 httpx==0.27.2
 httpx-sse==0.4.0
 huggingface-hub==0.26.2
@@ -28,6 +41,7 @@ jsonpatch==1.33
 jsonpointer==3.0.0
 jsonschema==4.23.0
 jsonschema-specifications==2024.10.1
+kiwisolver==1.4.7
 langchain==0.3.7
 langchain-community==0.3.7
 langchain-core==0.3.19
@@ -36,6 +50,7 @@ langsmith==0.1.143
 markdown-it-py==3.0.0
 MarkupSafe==3.0.2
 marshmallow==3.23.1
+matplotlib==3.9.3
 mdurl==0.1.2
 mpmath==1.3.0
 multidict==6.1.0
@@ -43,19 +58,26 @@ mypy-extensions==1.0.0
 narwhals==1.14.1
 networkx==3.4.2
 numpy==1.26.4
+opencv-python==4.10.0.84
 orjson==3.10.11
 packaging==24.2
 pandas==2.2.3
 pillow==11.0.0
 propcache==0.2.0
+proto-plus==1.25.0
 protobuf==5.28.3
+psutil==6.1.0
+py-cpuinfo==9.0.0
 pyarrow==18.0.0
+pyasn1==0.6.1
+pyasn1_modules==0.4.1
 pydantic==2.9.2
 pydantic-settings==2.6.1
 pydantic_core==2.23.4
 pydeck==0.9.1
 Pygments==2.18.0
 pymongo==4.10.1
+pyparsing==3.2.0
 python-dateutil==2.9.0.post0
 python-dotenv==1.0.1
 pytz==2024.2
@@ -66,7 +88,10 @@ requests==2.32.3
 requests-toolbelt==1.0.0
 rich==13.9.4
 rpds-py==0.21.0
+rsa==4.9
 safetensors==0.4.5
+scipy==1.14.1
+seaborn==0.13.2
 six==1.16.0
 smmap==5.0.1
 sniffio==1.3.1
@@ -77,11 +102,15 @@ tenacity==9.0.0
 tokenizers==0.20.3
 toml==0.10.2
 torch==2.5.1
+torchvision==0.20.1
 tornado==6.4.1
 tqdm==4.67.0
 transformers==4.46.3
 typing-inspect==0.9.0
 typing_extensions==4.12.2
 tzdata==2024.2
+ultralytics==8.3.47
+ultralytics-thop==2.0.12
+uritemplate==4.1.1
 urllib3==2.2.3
 yarl==1.17.2

yolo_model.py

@@ -0,0 +1,35 @@
+# YOLO model
+from ultralytics import YOLO
+# Streamlit
+import streamlit as st
+
+@st.cache_resource
+def load_yolo_model():
+    return YOLO("models/best.pt")
+
+def get_detected_objects(yolo_model, image_path, conf_threshold=0.5):
+    """
+    Run YOLO prediction on an image and return detected objects as a string.
+    
+    Parameters:
+        model_path (str): Path to the YOLO model file.
+        image_path (str): Path to the input image.
+        conf_threshold (float): Confidence threshold for detections.
+        
+    Returns:
+        str: A comma-separated string of detected object names.
+    """
+    # Load the YOLO model
+    model = yolo_model
+
+    # Run prediction
+    results = model.predict(source=image_path, conf=conf_threshold)
+
+    # Extract detected objects as a list
+    detected_objects = [box.cls for box in results[0].boxes]  # Access the first image's detections
+
+    # Convert class indices to class names
+    detected_class_names = [model.names[int(cls)] for cls in detected_objects]
+
+    # Join detected class names into a single string
+    return ", ".join(detected_class_names)