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()