app.py

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 supervision as sv
import torch
from PIL import Image
from tqdm import tqdm
import sys
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()
    
    # 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=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
            )
            
    # Examples section - Fixed version
    gr.Examples(
        examples=IMAGE_PROCESSING_EXAMPLES,
        inputs=[image_input, prompt_input],
        outputs=[output_gallery, analysis_output, detected_modality],
        fn=process_image,
        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()