app.py

import gradio as gr
import time
from collections import defaultdict
import spacy

# Load the symptom-to-disease mapping
symptom_data = {
    "Shortness of breath": {
        "questions": [
            "Do you also have chest pain?",
            "Do you feel fatigued often?",
            "Have you noticed swelling in your legs?"
        ],
        "diseases": ["Atelectasis", "Emphysema", "Edema"],
        "weights_yes": [30, 30, 40],
        "weights_no": [10, 20, 30]
    },
    "Persistent cough": {
        "questions": [
            "Is your cough dry or with mucus?",
            "Do you experience fever?",
            "Do you have difficulty breathing?"
        ],
        "diseases": ["Pneumonia", "Fibrosis", "Infiltration"],
        "weights_yes": [35, 30, 35],
        "weights_no": [10, 15, 20]
    },
    "Sharp chest pain": {
        "questions": [
            "Does it worsen with deep breaths?",
            "Do you feel lightheaded?",
            "Have you had recent trauma or surgery?"
        ],
        "diseases": ["Pneumothorax", "Effusion", "Cardiomegaly"],
        "weights_yes": [40, 30, 30],
        "weights_no": [15, 20, 25]
    },
    "Fatigue & swelling": {
        "questions": [
            "Do you feel breathless when lying down?",
            "Have you gained weight suddenly?",
            "Do you experience irregular heartbeat?"
        ],
        "diseases": ["Edema", "Cardiomegaly"],
        "weights_yes": [50, 30, 20],
        "weights_no": [20, 15, 15]
    },
    "Chronic wheezing": {
        "questions": [
            "Do you have a history of smoking?",
            "Do you feel tightness in your chest?",
            "Do you have frequent lung infections?"
        ],
        "diseases": ["Emphysema", "Fibrosis"],
        "weights_yes": [40, 30, 30],
        "weights_no": [15, 25, 20]
    }
}

# Load spaCy model for NLP
nlp = spacy.load("en_core_web_lg")

# Function to extract key symptom from user input
def extract_symptom(user_input):
    # Define the symptoms that the chatbot recognizes
    known_symptoms = list(symptom_data.keys())
    
    # Process the input with spaCy NLP model
    user_doc = nlp(user_input.lower())
    
    # Check if any of the known symptoms are in the user input
    for symptom in known_symptoms:
        if symptom.lower() in user_input.lower():
            return symptom
    
    # If no direct match, use similarity to find the closest symptom
    similarities = {}
    for symptom in known_symptoms:
        symptom_doc = nlp(symptom.lower())
        similarity = user_doc.similarity(symptom_doc)
        similarities[symptom] = similarity
    
    # Return the symptom with the highest similarity
    return max(similarities, key=similarities.get)

# Mapping of unrecognized symptoms to similar known ones
synonym_mapping = {
    "chest pain": "Sharp chest pain",
    "pain in chest": "Sharp chest pain",
    "wheezing": "Chronic wheezing",
    "cough": "Persistent cough",
    "shortness of breath": "Shortness of breath",
    "fatigue": "Fatigue & swelling"
}

# Global variables to track user state
user_state = {}

def chatbot(user_input):
    if "state" not in user_state:
        user_state["state"] = "greet"
    
    if user_state["state"] == "greet":
        user_state["state"] = "ask_symptom"
        return "Hello! I'm a medical AI assistant. Please describe your primary symptom."
    
    elif user_state["state"] == "ask_symptom":
        # Check if the symptom contains any synonym or keyword
        matched_symptom = None
        
        for synonym, recognized_symptom in synonym_mapping.items():
            if synonym in user_input.lower():
                matched_symptom = recognized_symptom
                break
        
        # If no synonym found, extract the symptom using NLP
        if not matched_symptom:
            matched_symptom = extract_symptom(user_input)
        
        # If the symptom is recognized, proceed to the next step
        if matched_symptom not in symptom_data:
            user_state["state"] = "ask_feeling"
            return "I'm sorry, I don't recognize that symptom. How do you feel?"
        
        user_state["symptom"] = matched_symptom
        user_state["state"] = "ask_duration"
        return "How long have you been experiencing this symptom? (Less than a week / More than a week)"
    
    elif user_state["state"] == "ask_feeling":
        # If the symptom is not recognized, ask how they feel
        return "Can you describe your symptoms in more detail?"
    
    elif user_state["state"] == "ask_duration":
        if user_input.lower() == "less than a week":
            user_state.clear()
            return "It might be a temporary issue. Please monitor your symptoms and consult a doctor if they persist."
        elif user_input.lower() == "more than a week":
            user_state["state"] = "follow_up"
            user_state["current_question"] = 0
            user_state["disease_scores"] = defaultdict(int)
            return symptom_data[user_state["symptom"]]["questions"][0]
        else:
            return "Please respond with 'Less than a week' or 'More than a week'."
    
    elif user_state["state"] == "follow_up":
        symptom = user_state["symptom"]
        question_index = user_state["current_question"]
        
        # Update probabilities
        if user_input.lower() == "yes":
            for i, disease in enumerate(symptom_data[symptom]["diseases"]):
                user_state["disease_scores"][disease] += symptom_data[symptom]["weights_yes"][i]
        else:
            for i, disease in enumerate(symptom_data[symptom]["diseases"]):
                user_state["disease_scores"][disease] += symptom_data[symptom]["weights_no"][i]
        
        # Move to the next question or finish
        user_state["current_question"] += 1
        if user_state["current_question"] < len(symptom_data[symptom]["questions"]):
            return symptom_data[symptom]["questions"][user_state["current_question"]]
        
        # Final diagnosis
        probable_disease = max(user_state["disease_scores"], key=user_state["disease_scores"].get)
        user_state.clear()
        return f"Based on your symptoms, the most likely condition is: {probable_disease}. Please consult a doctor for confirmation."

# Gradio Chatbot UI with improved features
with gr.Blocks() as demo:
    gr.Markdown("# Conversational Image Recognition Assistant: AI-Powered X-ray Diagnosis for Healthcare")
    chatbot_ui = gr.Chatbot()
    user_input = gr.Textbox(placeholder="Enter your response...", label="Your Message")
    submit = gr.Button("Send")
    clear_chat = gr.Button("Clear Chat")
    
    def respond(user_message, history):
        history.append((user_message, "Thinking..."))  # Show thinking message
        yield history, ""  # Immediate update
        
        time.sleep(1.5)  # Simulate processing delay
        bot_response = chatbot(user_message)
        history[-1] = (user_message, bot_response)  # Update with real response
        yield history, ""
    
    submit.click(respond, [user_input, chatbot_ui], [chatbot_ui, user_input])
    user_input.submit(respond, [user_input, chatbot_ui], [chatbot_ui, user_input])
    clear_chat.click(lambda: ([], ""), outputs=[chatbot_ui, user_input])

demo.launch()