This demo integrates BiomedParse, a foundation model for joint segmentation, detection, and recognition across 9 biomedical imaging modalities.
The model supports CT, MRI, X-Ray, Pathology, Ultrasound, Endoscope, Fundus, Dermoscopy, and OCT.
"""
IMAGE_PROCESSING_EXAMPLES = [
["BiomedParse Segmentation",
"https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/T0011.jpg",
"Optic disc in retinal Fundus"],
["BiomedParse Segmentation",
"https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/Part_3_226_pathology_breast.png",
"optic disc, optic cup"],
["BiomedParse Segmentation",
"https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/covid_1585.png",
"COVID-19 infection in chest X-Ray"],
["BiomedParse Segmentation",
"https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/TCGA_HT_7856_19950831_8_MRI-FLAIR_brain.png",
"Lower-grade glioma in brain MRI"],
["BiomedParse Segmentation",
"https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/LIDC-IDRI-0140_143_280_CT_lung.png",
"COVID-19 infection in chest CT"],
["BiomedParse Segmentation",
"https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/144DME_as_F.jpeg",
"Cystoid macular edema in retinal OCT"],
["BiomedParse Segmentation",
"https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/Part_1_516_pathology_breast.png",
"Glandular structure in colon Pathology"],
["BiomedParse Segmentation",
"https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/ISIC_0015551.jpg",
"Melanoma in skin Dermoscopy"],
["BiomedParse Segmentation",
"https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/C3_EndoCV2021_00462.jpg",
"Neoplastic polyp in colon Endoscope"]
]
BIOMEDPARSE_MODES = {
"CT-Abdomen": ["abdomen", "liver"],
"CT-Chest": ["lung"],
"CT-Liver": ["liver"],
"MRI-Abdomen": ["abdomen"],
"MRI-Cardiac": ["heart"],
"MRI-FLAIR-Brain": ["brain"],
"MRI-T1-Gd-Brain": ["brain"],
"Pathology": ["bladder", "breast", "cervix", "colon", "esophagus", "kidney",
"liver", "ovarian", "prostate", "stomach", "testis", "thyroid", "uterus"],
"X-Ray-Chest": ["chest"],
"Ultrasound-Cardiac": ["heart"],
"Endoscopy": ["colon"],
"Fundus": ["retinal"],
"Dermoscopy": ["skin"],
"OCT": ["retinal"]
}
IMAGE_INFERENCE_MODES = [
"BIOMED SEGMENTATION",
"BIOMED DETECTION",
"BIOMED RECOGNITION",
"BIOMED SEGMENTATION + DETECTION",
"BIOMED SEGMENTATION + RECOGNITION",
"BIOMED DETECTION + RECOGNITION",
"BIOMED SEGMENTATION + DETECTION + RECOGNITION"
]
MODALITY_PROMPTS = {
"CT-Abdomen": ["postcava", "aorta", "right kidney", "kidney", "left kidney", "duodenum", "pancreas", "liver", "spleen", "stomach", "gallbladder", "left adrenal gland", "adrenal gland", "right adrenal gland", "esophagus"],
"CT-Chest": ["nodule", "COVID-19 infection", "tumor"],
"MRI-Abdomen": ["aorta", "postcava", "right kidney", "duodenum", "kidney", "left kidney", "liver", "pancreas", "gallbladder", "stomach", "spleen", "left adrenal gland", "adrenal gland", "right adrenal gland", "esophagus"],
"MRI-Cardiac": ["left heart ventricle", "myocardium", "right heart ventricle"],
"MRI-FLAIR-Brain": ["edema", "tumor core", "whole tumor"],
"MRI-T1-Gd-Brain": ["enhancing tumor", "non-enhancing tumor", "tumor core"],
"Pathology": ["connective tissue cells", "inflammatory cells", "neoplastic cells", "epithelial cells"],
"X-Ray-Chest": ["left lung", "lung", "right lung"],
"Ultrasound-Cardiac": ["left heart atrium", "left heart ventricle"],
"Endoscopy": ["neoplastic polyp", "polyp", "non-neoplastic polyp"],
"Fundus": ["optic cup", "optic disc"],
"Dermoscopy": ["lesion", "melanoma"],
"OCT": ["edema"] }
def extract_modality_from_llm(llm_output):
"""Extract modality from LLM output and map it to BIOMEDPARSE_MODES"""
llm_output = llm_output.lower()
# Direct modality mapping
modality_keywords = {
'ct': {
'abdomen': 'CT-Abdomen',
'chest': 'CT-Chest',
'liver': 'CT-Liver'
},
'mri': {
'abdomen': 'MRI-Abdomen',
'cardiac': 'MRI-Cardiac',
'heart': 'MRI-Cardiac',
'flair': 'MRI-FLAIR-Brain',
't1': 'MRI-T1-Gd-Brain',
'contrast': 'MRI-T1-Gd-Brain',
'brain': 'MRI-FLAIR-Brain' # default to FLAIR if just "brain" is mentioned
},
'x-ray': {'chest': 'X-Ray-Chest'},
'ultrasound': {'cardiac': 'Ultrasound-Cardiac', 'heart': 'Ultrasound-Cardiac'},
'endoscopy': {'': 'Endoscopy'},
'fundus': {'': 'Fundus'},
'dermoscopy': {'': 'Dermoscopy'},
'oct': {'': 'OCT'},
'pathology': {'': 'Pathology'}
}
for modality, subtypes in modality_keywords.items():
if modality in llm_output:
# For modalities with subtypes, try to find the specific subtype
if subtypes:
for keyword, specific_modality in subtypes.items():
if not keyword or keyword in llm_output:
return specific_modality
# For modalities without subtypes, return the direct mapping
return next(iter(subtypes.values()))
return None
def extract_clinical_findings(llm_output, modality):
"""Extract relevant clinical findings that match available anatomical sites in BIOMEDPARSE_MODES"""
available_sites = BIOMEDPARSE_MODES.get(modality, [])
findings = []
# Convert sites to lowercase for case-insensitive matching
sites_lower = {site.lower(): site for site in available_sites}
# Look for each available site in the LLM output
for site_lower, original_site in sites_lower.items():
if site_lower in llm_output.lower():
findings.append(original_site)
# Add additional findings from MODALITY_PROMPTS if available
if modality in MODALITY_PROMPTS:
for prompt in MODALITY_PROMPTS[modality]:
if prompt.lower() in llm_output.lower() and prompt not in findings:
findings.append(prompt)
return findings
def on_mode_dropdown_change(selected_mode):
if selected_mode in IMAGE_INFERENCE_MODES:
# Show modality dropdown and hide other inputs initially
return [
gr.Dropdown(visible=True, choices=list(BIOMEDPARSE_MODES.keys()), label="Modality"),
gr.Dropdown(visible=True, label="Anatomical Site"),
gr.Textbox(visible=False),
gr.Textbox(visible=False)
]
else:
# Original behavior for other modes
return [
gr.Dropdown(visible=False),
gr.Dropdown(visible=False),
gr.Textbox(visible=True),
gr.Textbox(visible=(selected_mode == None))
]
def on_modality_change(modality):
if modality:
return gr.Dropdown(choices=BIOMEDPARSE_MODES[modality], visible=True)
return gr.Dropdown(visible=False)
def initialize_model():
opt = load_opt_from_config_files(["configs/biomedparse_inference.yaml"])
pretrained_pth = 'hf_hub:microsoft/BiomedParse'
opt['device'] = 'cuda' if torch.cuda.is_available() else 'cpu'
model = BaseModel(opt, build_model(opt)).from_pretrained(pretrained_pth).eval()
with torch.no_grad():
model.model.sem_seg_head.predictor.lang_encoder.get_text_embeddings(
BIOMED_CLASSES + ["background"], is_eval=True
)
return model
def initialize_llm():
try:
print("Starting LLM initialization...")
# Add quantization config
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16
)
model = AutoModel.from_pretrained(
"ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
device_map="auto",
torch_dtype=torch.float16,
trust_remote_code=True,
low_cpu_mem_usage=True,
quantization_config=quantization_config
)
print("Model loaded successfully")
tokenizer = AutoTokenizer.from_pretrained(
"ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
trust_remote_code=True
)
print("Tokenizer loaded successfully")
return model, tokenizer
except Exception as e:
print(f"Failed to initialize LLM: {str(e)}")
return None, None
model = initialize_model()
llm_model, llm_tokenizer = initialize_llm()
def update_example_prompts(modality):
if modality in MODALITY_PROMPTS:
examples = MODALITY_PROMPTS[modality]
return f"Example prompts for {modality}:\n" + ", ".join(examples)
return ""
# Utility functions
@spaces.GPU
@torch.inference_mode()
@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
def process_image(image_path, user_prompt, modality=None):
try:
if not image_path:
raise ValueError("Please upload an image")
image = read_rgb(image_path)
pil_image = Image.fromarray(image)
# Step 1: Get LLM analysis
question = f"Analyze this medical image considering the following context: {user_prompt}. Include modality, anatomical structures, and any abnormalities."
msgs = [{'role': 'user', 'content': [pil_image, question]}]
llm_response = ""
for new_text in llm_model.chat(
image=pil_image,
msgs=msgs,
tokenizer=llm_tokenizer,
sampling=True,
temperature=0.95,
stream=True
):
llm_response += new_text
# Step 2: Extract modality from LLM output
detected_modality = extract_modality_from_llm(llm_response)
if not detected_modality:
raise ValueError("Could not determine image modality from LLM output")
# Step 3: Extract relevant clinical findings
clinical_findings = extract_clinical_findings(llm_response, detected_modality)
# Step 4: Generate masks for each finding
results = []
analysis_results = []
colors = [(255,0,0), (0,255,0), (0,0,255), (255,255,0), (255,0,255)] # Different colors for different findings
for idx, finding in enumerate(clinical_findings):
pred_mask = interactive_infer_image(model, pil_image, [finding])[0]
p_value = check_mask_stats(image, pred_mask * 255, detected_modality, finding)
analysis_results.append(f"P-value for '{finding}' ({detected_modality}): {p_value:.4f}")
# Create colored overlay
overlay_image = image.copy()
color = colors[idx % len(colors)]
overlay_image[pred_mask > 0.5] = color
results.append(overlay_image)
# Update LLM response with color references
enhanced_response = llm_response + "\n\nSegmentation Results:\n"
for idx, finding in enumerate(clinical_findings):
color_name = ["red", "green", "blue", "yellow", "magenta"][idx % len(colors)]
enhanced_response += f"- {finding} (shown in {color_name})\n"
combined_analysis = "\n\n" + "="*50 + "\n"
combined_analysis += "BiomedParse Analysis:\n"
combined_analysis += "\n".join(analysis_results)
combined_analysis += "\n\n" + "="*50 + "\n"
combined_analysis += "Enhanced LLM Analysis:\n"
combined_analysis += enhanced_response
combined_analysis += "\n" + "="*50
return results, combined_analysis, detected_modality
except Exception as e:
error_msg = f"⚠️ An error occurred: {str(e)}"
print(f"Error details: {str(e)}", flush=True)
return None, error_msg
# Define Gradio interface
with gr.Blocks() as demo:
gr.HTML(MARKDOWN)
with gr.Row():
with gr.Column():
image_input = gr.Image(type="filepath", label="Input Image")
prompt_input = gr.Textbox(
lines=4,
placeholder="Ask any question about the medical image...",
label="Question/Prompt"
)
detected_modality = gr.Textbox(
label="Detected Modality",
interactive=False,
visible=True
)
submit_btn = gr.Button("Analyze")
with gr.Column():
output_gallery = gr.Gallery(
label="Segmentation Results",
show_label=True,
columns=[2],
height="auto"
)
analysis_output = gr.Textbox(
label="Analysis",
interactive=False,
show_label=True,
lines=10
)
# Examples section
gr.Examples(
examples=IMAGE_PROCESSING_EXAMPLES,
inputs=[image_input, prompt_input],
outputs=[output_gallery, analysis_output, detected_modality],
cache_examples=True,
)
# Connect the submit button to the process_image function
submit_btn.click(
fn=process_image,
inputs=[image_input, prompt_input],
outputs=[output_gallery, analysis_output, detected_modality],
api_name="process"
)
demo.launch()
