import streamlit as st import tensorflow as tf import os # Import your utility functions from utils import ( predict_multi_line_text, tokenizer, ) from config import index_to_label, acronyms_to_entities, MAX_LENGTH from metrics import precision, recall, f1_score # Register the custom metric functions tf.keras.utils.get_custom_objects()[precision.__name__] = precision tf.keras.utils.get_custom_objects()[recall.__name__] = recall tf.keras.utils.get_custom_objects()[f1_score.__name__] = f1_score # Load your trained model model_dir = './model' # Adjust the path as needed model_1 = tf.keras.models.load_model(os.path.join(model_dir, 'model_1.h5')) # Define label colors for different entity types suitable for dark background LABEL_COLORS = { 'Activity': '#FF7F50', # Coral 'Administration': '#6495ED', # Cornflower Blue 'Age': '#FFB6C1', # Light Pink 'Area': '#7FFF00', # Chartreuse 'Biological_attribute': '#FFD700', # Gold 'Biological_structure': '#00FA9A', # Medium Spring Green 'Clinical_event': '#BA55D3', # Medium Orchid 'Color': '#00CED1', # Dark Turquoise 'Coreference': '#FFA07A', # Light Salmon 'Date': '#ADFF2F', # Green Yellow 'Detailed_description': '#DA70D6', # Orchid 'Diagnostic_procedure': '#87CEFA', # Light Sky Blue 'Disease_disorder': '#FF4500', # Orange Red 'Distance': '#32CD32', # Lime Green 'Dosage': '#8A2BE2', # Blue Violet 'Duration': '#F08080', # Light Coral 'Family_history': '#20B2AA', # Light Sea Green 'Frequency': '#FF6347', # Tomato 'Height': '#4682B4', # Steel Blue 'History': '#EE82EE', # Violet 'Lab_value': '#FFDAB9', # Peach Puff 'Mass': '#7B68EE', # Medium Slate Blue 'Medication': '#00FF7F', # Spring Green 'Nonbiological_location': '#FF69B4', # Hot Pink 'Occupation': '#BDB76B', # Dark Khaki 'Other_entity': '#D3D3D3', # Light Grey 'Other_event': '#FF1493', # Deep Pink 'Outcome': '#00BFFF', # Deep Sky Blue 'Personal_background': '#00FFFF', # Aqua 'Qualitative_concept': '#FFA500', # Orange 'Quantitative_concept': '#FFA500', # Orange (same as above) 'Severity': '#1E90FF', # Dodger Blue 'Sex': '#FF00FF', # Magenta 'Shape': '#40E0D0', # Turquoise 'Sign_symptom': '#FF69B4', # Hot Pink 'Subject': '#F0E68C', # Khaki 'Texture': '#98FB98', # Pale Green 'Therapeutic_procedure': '#8B008B', # Dark Magenta 'Time': '#DC143C', # Crimson 'Volume': '#5F9EA0', # Cadet Blue 'Weight': '#FA8072', # Salmon } # Define the prediction function def predict_ner(text): try: # Predict entities entities = predict_multi_line_text( text, model_1, index_to_label, acronyms_to_entities, MAX_LENGTH ) # Sort entities by their start position entities = sorted(entities, key=lambda x: x[0]) # Build HTML string with highlighted entities html_output = "" last_idx = 0 for start, end, label in entities: # Append text before the entity if last_idx < start: html_output += text[last_idx:start] # Get the color for the label, default to light grey if not specified color = LABEL_COLORS.get(label, '#D3D3D3') # Light grey # Wrap the entity with a span tag including style entity_text = text[start:end] # Include the label next to the entity html_output += f'''{entity_text} [{label}]''' last_idx = end # Append any remaining text if last_idx < len(text): html_output += text[last_idx:] return html_output except Exception as e: return f"
Error: {str(e)}
" # Set up the Streamlit app with dark theme st.set_page_config(page_title="Medical NER", page_icon="🩺", layout="wide") # Apply custom CSS for dark background and text colors st.markdown( """ """, unsafe_allow_html=True ) st.title("🩺 Medical Named Entity Recognition") st.markdown(""" Enter medical text below to identify and highlight entities such as diseases, medications, and anatomical terms. """) # Input text area text_input = st.text_area("Enter medical text here:", height=200) # Analyze button if st.button("Analyze"): if text_input.strip(): with st.spinner("Analyzing..."): result = predict_ner(text_input) # Display the result with HTML rendering st.markdown(f"