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_and_prompts(llm_output):
"""
Extract modality and relevant prompts from LLM output
Returns: (modality_type, list_of_prompts)
"""
llm_output = llm_output.lower()
# Dictionary mapping keywords to modalities
modality_indicators = {
'dermatoscop': 'Dermoscopy',
'dermatoscope': 'Dermoscopy',
'dermal': 'Dermoscopy',
'skin lesion': 'Dermoscopy',
'dermatological': 'Dermoscopy',
'oct': 'OCT',
'optical coherence': 'OCT',
'fundus': 'Fundus',
'retina': 'Fundus',
'endoscop': 'Endoscopy',
'colon': 'Endoscopy',
'pathological': 'Pathology',
# 'tissue': 'Pathology',
'histolog': 'Pathology',
'x-ray': 'X-Ray-Chest',
'xray': 'X-Ray-Chest',
'chest radiograph': 'X-Ray-Chest',
'mri': None, # Will be refined below
'magnetic resonance': None, # Will be refined below
'ct': None, # Will be refined below
'computed tomography': None, # Will be refined below
'ultrasound': 'Ultrasound-Cardiac',
'sonograph': 'Ultrasound-Cardiac'
}
# First pass: Detect base modality
detected_modality = None
for keyword, modality in modality_indicators.items():
if keyword in llm_output:
detected_modality = modality
break
# Second pass: Refine MRI and CT if detected
if detected_modality is None and ('mri' in llm_output or 'magnetic resonance' in llm_output):
if 'brain' in llm_output or 'flair' in llm_output:
detected_modality = 'MRI-FLAIR-Brain'
elif 'cardiac' in llm_output or 'heart' in llm_output:
detected_modality = 'MRI-Cardiac'
elif 'abdomen' in llm_output:
detected_modality = 'MRI-Abdomen'
elif 't1' in llm_output or 'contrast' in llm_output:
detected_modality = 'MRI-T1-Gd-Brain'
else:
detected_modality = 'MRI'
if detected_modality is None and ('ct' in llm_output or 'computed tomography' in llm_output):
if 'chest' in llm_output or 'lung' in llm_output:
detected_modality = 'CT-Chest'
elif 'liver' in llm_output:
detected_modality = 'CT-Liver'
elif 'abdomen' in llm_output:
detected_modality = 'CT-Abdomen'
else:
detected_modality = 'CT'
# If still no modality detected, return None
if not detected_modality:
return "", []
# Get relevant prompts for the detected modality
if detected_modality in MODALITY_PROMPTS:
relevant_prompts = MODALITY_PROMPTS[detected_modality]
else:
relevant_prompts = []
return detected_modality, relevant_prompts
def on_mode_dropdown_change(selected_mode):
if selected_mode in IMAGE_INFERENCE_MODES:
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:
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...")
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 ""
@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:
return [], "Please upload an image", "No modality detected"
image = read_rgb(image_path)
pil_image = Image.fromarray(image)
# question = (
# f"Analyze this medical image considering the following context: {user_prompt}. "
# "Include modality, anatomical structures, and any abnormalities."
# )
question = 'What type of medical imaging modality is this? which organ? Be specific.'
msgs = [{'role': 'user', 'content': [pil_image, question]}]
llm_response = ""
if llm_model and llm_tokenizer:
try:
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
except Exception as e:
print(f"LLM chat error: {str(e)}")
llm_response = "LLM analysis failed. Proceeding with basic analysis."
else:
llm_response = "LLM not available. Please check LLM initialization logs."
detected_modality, relevant_prompts = extract_modality_and_prompts(llm_response)
if not detected_modality:
detected_modality = "X-Ray-Chest" # Fallback modality
relevant_prompts = MODALITY_PROMPTS["X-Ray-Chest"]
# results = []
# analysis_results = []
# colors = [(255,0,0), (0,255,0), (0,0,255), (255,255,0), (255,0,255)]
# # Add color mapping to analysis with more natural language
# color_descriptions = []
# for idx, prompt in enumerate(relevant_prompts):
# color = colors[idx % len(colors)]
# color_name = {(255,0,0): "red", (0,255,0): "green", (0,0,255): "blue",
# (255,255,0): "yellow", (255,0,255): "magenta"}[color]
# color_descriptions.append(f"The {prompt} is highlighted in {color_name} color")
# for idx, prompt in enumerate(relevant_prompts):
# try:
# mask_list = interactive_infer_image(model, pil_image, [prompt])
# if mask_list is None or len(mask_list) == 0:
# analysis_results.append(f"No mask generated for '{prompt}'")
# continue
# pred_mask = mask_list[0]
# # Check if mask is valid using numpy's any() function
# if pred_mask is None or not np.any(pred_mask):
# analysis_results.append(f"Empty mask generated for '{prompt}'")
# continue
# overlay_image = image.copy()
# color = colors[idx % len(colors)]
# mask_indices = pred_mask > 0.5
# if np.any(mask_indices): # Use np.any() for boolean array check
# overlay_image[mask_indices] = color
# results.append(overlay_image)
# except Exception as e:
# print(f"Error processing finding {prompt}: {str(e)}")
# analysis_results.append(f"Failed to process '{prompt}': {str(e)}")
# if not results:
# results = [image] # Return original image if no overlays were created
detailed_analysis = ""
if llm_model and llm_tokenizer:
try:
# Add color legend with more natural language
# detailed_analysis += "\n\n As shown in the images outputs details:\n \n" + "\n".join(color_descriptions)
analysis_prompt = f"Focus more on the user question. which is: {user_prompt}. Give the modality, organ, analysis, abnormalities (if any), treatment (if abnormalities are present) for this image. "
msgs = [{'role': 'user', 'content': [pil_image, analysis_prompt]}]
for new_text in llm_model.chat(
image=pil_image,
msgs=msgs,
tokenizer=llm_tokenizer,
sampling=True,
temperature=0.95,
stream=True
):
detailed_analysis += new_text
# Add segmentation details in a more natural way
results = []
analysis_results = []
colors = [(255,0,0), (0,255,0), (0,0,255), (255,255,0), (255,0,255)]
# Add color mapping to analysis with more natural language
color_descriptions = []
# First loop: collect prompts found in the analysis
found_prompts = []
count = 0
for idx, prompt in enumerate(relevant_prompts):
if prompt in detailed_analysis.lower():
color = colors[count % len(colors)]
color_name = {(255,0,0): "red", (0,255,0): "green", (0,0,255): "blue", (255,255,0): "yellow", (255,0,255): "magenta"}[color]
color_descriptions.append(f"The {prompt} is highlighted in {color_name} color for reference")
found_prompts.append(prompt)
count += 1
# Second loop: only process prompts found in analysis
for idx, prompt in enumerate(found_prompts):
try:
mask_list = interactive_infer_image(model, pil_image, [prompt])
if mask_list is None or len(mask_list) == 0:
analysis_results.append(f"No mask generated for '{prompt}'")
continue
pred_mask = mask_list[0]
# Check if mask is valid using numpy's any() function
if pred_mask is None or not np.any(pred_mask):
analysis_results.append(f"Empty mask generated for '{prompt}'")
continue
overlay_image = image.copy()
color = colors[idx % len(colors)]
mask_indices = pred_mask > 0.5
if np.any(mask_indices): # Use np.any() for boolean array check
overlay_image[mask_indices] = color
results.append(overlay_image)
except Exception as e:
print(f"Error processing finding {prompt}: {str(e)}")
analysis_results.append(f"Failed to process '{prompt}': {str(e)}")
if not results:
results = [image] # Return original image if no overlays were created
detailed_analysis += ""
if color_descriptions:
detailed_analysis += " " + " ".join(color_descriptions) + "."
else:
detailed_analysis += " No significant segments were detected."
except Exception as e:
print(f"LLM chat error: {str(e)}")
detailed_analysis = "LLM analysis failed. Proceeding with basic analysis."
if color_descriptions:
detailed_analysis += "\n\nHowever, in the segmentation analysis: " + " ".join(color_descriptions) + "."
else:
detailed_analysis = "LLM not available. Please check LLM initialization logs."
if color_descriptions:
detailed_analysis += "\n\nIn the segmentation analysis: " + " ".join(color_descriptions) + "."
return results, detailed_analysis, detected_modality
except Exception as e:
error_msg = f"⚠️ An error occurred: {str(e)}"
print(f"Error details: {str(e)}", flush=True)
return [image] if 'image' in locals() else [], error_msg, "Error detecting modality"
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=2,
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
)
# Add this to clear outputs when input image is cleared
image_input.clear(
lambda: ([], "", ""),
outputs=[output_gallery, analysis_output, detected_modality]
)
submit_btn.click(
fn=process_image,
inputs=[image_input, prompt_input],
outputs=[output_gallery, analysis_output, detected_modality],
api_name="process"
)
demo.launch()
