|
import gradio as gr |
|
import os |
|
from typing import Tuple, Optional |
|
import os |
|
import shutil |
|
import sys |
|
from pathlib import Path |
|
import cv2 |
|
import gradio as gr |
|
import numpy as np |
|
import spaces |
|
import torch |
|
from PIL import Image |
|
from tqdm import tqdm |
|
from pathlib import Path |
|
from huggingface_hub import login |
|
from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig |
|
|
|
token = os.getenv("HF_TOKEN") |
|
if token: |
|
login(token=token) |
|
current_dir = Path(__file__).parent |
|
sys.path.append(str(current_dir)) |
|
|
|
from modeling.BaseModel import BaseModel |
|
from modeling import build_model |
|
from utilities.arguments import load_opt_from_config_files |
|
from utilities.constants import BIOMED_CLASSES |
|
from inference_utils.inference import interactive_infer_image |
|
from inference_utils.output_processing import check_mask_stats |
|
from inference_utils.processing_utils import read_rgb |
|
|
|
import spaces |
|
|
|
|
|
MARKDOWN = """ |
|
<div align="center" style="padding: 20px 0;"> |
|
<h1 style="font-size: 3em; margin: 0;"> |
|
ሀ<span style="color: #32CD32;">A</span>ኪ<span style="color: #FFD700;">i</span>ም |
|
<sup style="font-size: 0.5em;">AI</sup> |
|
</h1> |
|
|
|
<div style="display: flex; justify-content: center; align-items: center; gap: 15px; margin: 15px 0;"> |
|
<a href="https://cyberbrainai.com/"> |
|
<img src="https://cyberbrainai.com/assets/logo.svg" alt="CyberBrain AI" style="width:40px; height:40px; vertical-align: middle;"> |
|
</a> |
|
<a href="https://colab.research.google.com/drive/1p3Yf_6xdZPMz5RUtt_NyxrDjrbSgvTDy#scrollTo=t30NqIrCKdAI"> |
|
<img src="https://colab.research.google.com/assets/colab-badge.svg" alt="ድinቅneሽ" style="vertical-align: middle;"> |
|
</a> |
|
<a href="https://www.youtube.com/watch?v=Dv003fTyO-Y"> |
|
<img src="https://badges.aleen42.com/src/youtube.svg" alt="YouTube" style="vertical-align: middle;"> |
|
</a> |
|
</div> |
|
</div> |
|
<div> |
|
<p style="font-size: 1.4em; line-height: 1.5; margin: 15px 0; text-align: left;"> |
|
This demo integrates BiomedParse, a foundation model for joint segmentation, detection, and recognition across 9 biomedical imaging modalities. |
|
The model supports <span style="color: #FF4500;">CT</span>, <span style="color: #4169E1;">MRI</span>, <span style="color: #32CD32;">X-Ray</span>, <span style="color: #9370DB;">Pathology</span>, <span style="color: #FFD700;">Ultrasound</span>, <span style="color: #FF69B4;">Endoscope</span>, <span style="color: #20B2AA;">Fundus</span>, <span style="color: #FF8C00;">Dermoscopy</span>, and <span style="color: #8B008B;">OCT</span>. |
|
</p> |
|
</div> |
|
|
|
""" |
|
|
|
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() |
|
|
|
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' |
|
}, |
|
'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 keyword, specific_modality in subtypes.items(): |
|
if not keyword or keyword in llm_output: |
|
return specific_modality |
|
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 = [] |
|
|
|
|
|
sites_lower = {site.lower(): site for site in available_sites} |
|
|
|
|
|
for site_lower, original_site in sites_lower.items(): |
|
if site_lower in llm_output.lower(): |
|
findings.append(original_site) |
|
|
|
|
|
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: |
|
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: |
|
raise ValueError("Please upload an image") |
|
|
|
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." |
|
) |
|
msgs = [{'role': 'user', 'content': [pil_image, question]}] |
|
|
|
llm_response = "" |
|
if llm_model and llm_tokenizer: |
|
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 |
|
else: |
|
llm_response = "LLM not available. Please check LLM initialization logs." |
|
|
|
detected_modality = extract_modality_from_llm(llm_response) |
|
if not detected_modality: |
|
|
|
detected_modality = "X-Ray-Chest" |
|
|
|
clinical_findings = extract_clinical_findings(llm_response, detected_modality) |
|
if not clinical_findings: |
|
|
|
clinical_findings = [detected_modality.split("-")[-1]] |
|
|
|
results = [] |
|
analysis_results = [] |
|
colors = [(255,0,0), (0,255,0), (0,0,255), (255,255,0), (255,0,255)] |
|
|
|
for idx, finding in enumerate(clinical_findings): |
|
mask_list = interactive_infer_image(model, pil_image, [finding]) |
|
if not mask_list: |
|
analysis_results.append(f"No mask found for '{finding}'.") |
|
continue |
|
|
|
pred_mask = mask_list[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}") |
|
|
|
overlay_image = image.copy() |
|
color = colors[idx % len(colors)] |
|
overlay_image[pred_mask > 0.5] = color |
|
results.append(overlay_image) |
|
|
|
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 |
|
|
|
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 |
|
) |
|
|
|
submit_btn.click( |
|
fn=process_image, |
|
inputs=[image_input, prompt_input], |
|
outputs=[output_gallery, analysis_output, detected_modality], |
|
api_name="process" |
|
) |
|
|
|
demo.launch() |
