app.py

# Imports standard
import torch
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import gradio as gr
import os
import subprocess
import sys

# Installation des dépendances nécessaires
subprocess.run(['apt-get', 'update'], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
packages = ['openmpi-bin', 'libopenmpi-dev']
command = ['apt-get', 'install', '-y'] + packages
subprocess.run(command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'mpi4py'])
subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'pydicom'])
subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'SimpleITK'])

# Imports Hugging Face
from huggingface_hub import hf_hub_download, login
import spaces
import numpy as np
from PIL import Image
import torch
import torch.nn.functional as F
from torchvision import transforms
from translator import translate_text

# Imports locaux
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 modeling.language.loss import vl_similarity
from inference_utils.output_processing import check_mask_stats

t = []
t.append(transforms.Resize((1024, 1024), interpolation=Image.BICUBIC))
transform = transforms.Compose(t)

def init_huggingface():
    """Initialize Hugging Face connection and download the model."""
    hf_token = os.getenv('HF_TOKEN')
    if hf_token is None:
        raise ValueError("Hugging Face token not found. Please set the HF_TOKEN environment variable.")
    login(hf_token)
  
    pretrained_path = hf_hub_download(
        repo_id="microsoft/BiomedParse",
        filename="biomedparse_v1.pt",
        local_dir="pretrained"
    )
    return pretrained_path

@spaces.GPU
def init_distributed(opt):
    opt['CUDA'] = opt.get('CUDA', True) and torch.cuda.is_available()
    if 'OMPI_COMM_WORLD_SIZE' not in os.environ:
        # application was started without MPI
        # default to single node with single process
        opt['env_info'] = 'no MPI'
        opt['world_size'] = 1
        opt['local_size'] = 1
        opt['rank'] = 0
        opt['local_rank'] = 0
        opt['master_address'] = '127.0.0.1'
        opt['master_port'] = '8673'
    else:
        # application was started with MPI
        # get MPI parameters
        opt['world_size'] = int(os.environ['OMPI_COMM_WORLD_SIZE'])
        opt['local_size'] = int(os.environ['OMPI_COMM_WORLD_LOCAL_SIZE'])
        opt['rank'] = int(os.environ['OMPI_COMM_WORLD_RANK'])
        opt['local_rank'] = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])

    # set up device
    if not opt['CUDA']:
        assert opt['world_size'] == 1, 'multi-GPU training without CUDA is not supported since we use NCCL as communication backend'
        opt['device'] = torch.device("cpu")
    else:
        torch.cuda.set_device(opt['local_rank'])
        opt['device'] = torch.device("cuda", opt['local_rank'])

    return opt

def setup_model():
    """Initialize the model on CPU without CUDA initialization."""
    opt = load_opt_from_config_files(["configs/biomedparse_inference.yaml"])
    opt = init_distributed(opt)
    opt['device'] = 'cuda' if torch.cuda.is_available() else 'cpu'
    pretrained_path = init_huggingface()
    model = BaseModel(opt, build_model(opt)).from_pretrained(pretrained_path).eval()
    return model

@torch.no_grad()
@spaces.GPU 
def predict_image(model, image, prompts):
    model = model.cuda()
    print("====================== Model moved to GPU via predict_image ======================")
    prompts = [translate_text(p, "fr", "en") for p in prompts]
    # Initialiser les embeddings textuels avant l'évaluation
    lang_encoder = model.model.sem_seg_head.predictor.lang_encoder
    lang_encoder.get_text_embeddings(BIOMED_CLASSES + ["background"], is_eval=True)   
    model.model.sem_seg_head.predictor.lang_encoder.get_text_embeddings(BIOMED_CLASSES + ["background"], is_eval=True)  
    # Préparer l'image
    image_resize = transform(image)
    width = image.size[0]
    height = image.size[1]
    image_resize = np.asarray(image_resize)
    image = torch.from_numpy(image_resize.copy()).permute(2,0,1).cuda()

    # Préparer les données d'entrée
    data = {"image": image, 'text': prompts, "height": height, "width": width}
    
    # Configurer les tâches
    model.model.task_switch['spatial'] = False
    model.model.task_switch['visual'] = False
    model.model.task_switch['grounding'] = True
    model.model.task_switch['audio'] = False

    # Évaluation
    batch_inputs = [data]
    results, image_size, extra = model.model.evaluate_demo(batch_inputs)

    # Traitement des prédictions
    pred_masks = results['pred_masks'][0]
    v_emb = results['pred_captions'][0]
    t_emb = extra['grounding_class']

    # Normalisation
    t_emb = t_emb / (t_emb.norm(dim=-1, keepdim=True) + 1e-7)
    v_emb = v_emb / (v_emb.norm(dim=-1, keepdim=True) + 1e-7)

    # Calcul de similarité
    temperature = lang_encoder.logit_scale
    out_prob = vl_similarity(v_emb, t_emb, temperature=temperature)
    
    # Sélection des masques
    matched_id = out_prob.max(0)[1]
    pred_masks_pos = pred_masks[matched_id,:,:]

    # Redimensionnement à la taille d'origine
    pred_mask_prob = F.interpolate(pred_masks_pos[None,], (data['height'], data['width']), 
                                   mode='bilinear')[0,:,:data['height'],:data['width']].sigmoid().cpu().detach().numpy()
    return pred_mask_prob


def process_image(image, prompts, model):
    """Process image with proper error handling."""
    try:
        if isinstance(image, str):
            image = Image.open(image)
        else:
            image = Image.fromarray(image)

        prompts = [p.strip() for p in prompts.split(',')]

        if not prompts:
            raise ValueError("No valid prompts provided")
            
        pred_masks = predict_image(model, image, prompts)
        fig = plt.figure(figsize=(10, 5))
        plt.subplot(1, len(pred_masks) + 1, 1)
        plt.imshow(image)
        plt.title('Image originale')
        plt.axis('off')

        for i, mask in enumerate(pred_masks):
            plt.subplot(1, len(pred_masks) + 1, i + 2)
            plt.imshow(image)
            plt.imshow(mask, alpha=0.5, cmap='Reds')
            plt.title(prompts[i])
            plt.axis('off')
        return fig
        
    except Exception as e:
        print(f"Error in process_image: {str(e)}")
        raise

def setup_gradio_interface(model):
    """Configure l'interface Gradio."""
    return gr.Interface(
        theme=gr.Theme.from_hub("JohnSmith9982/small_and_pretty"),
        fn=lambda img, txt: process_image(img, txt, model),
        inputs=[
            gr.Image(type="numpy", label="Image médicale"),
            gr.Textbox(
                label="Prompts (séparés par des virgules)",
                placeholder="edema, lesion, etc...",
                elem_classes="white"
            )
        ],
        outputs=gr.Plot(),
        title=" 🇬🇦 Core IA - Traitement d'image medicale",
        description="""Chargez une image médicale de type (IRM , Echographie,  ) et spécifiez les éléments à segmenter : Les cas d’utilisation incluent des tâches variées d’analyse d’images médicales. En imagerie CT, le modèle peut détecter et segmenter des organes et pathologies dans des régions telles que l’abdomen, le côlon, le foie, les poumons ou le bassin. Avec l’IRM, il est capable de traiter des structures abdominales, cérébrales, cardiaques et prostatiques selon les différentes séquences (FLAIR, T1-Gd, T2). En radiographie, il peut identifier des anomalies pulmonaires ou des infections liées à la COVID-19.

Le modèle couvre également d’autres modalités comme la dermoscopie, l’endoscopie, le fond d’œil, et la pathologie, avec des applications pour la détection de lésions, de polypes ou de cellules néoplasiques dans divers tissus et organes. En échographie, il identifie des anomalies mammaires, cardiaques ou fœtales. Enfin, en TOCT, il peut analyser des structures rétiniennes.

Bien que performant dans ces contextes, il est important de considérer les spécificités des ensembles de données externes et de procéder à un ajustement pour des résultats précis.""",
        examples=[
            ["examples/144DME_as_F.jpeg", "Dans cette image donne moi l'œdème"],
            ["examples/T0011.jpg", "disque optique, cupule optique"],
            ["examples/C3_EndoCV2021_00462.jpg", "Trouve moi le polyp"],
            ["examples/covid_1585.png", "Qu'est ce qui ne va pas ici ?"],
            ['examples/Part_1_516_pathology_breast.png', "cellules néoplasiques , cellules inflammatoires ,  cellules du tissu conjonctif"]
        ]
    )

def main():
    """Entry point avoiding CUDA initialization in main process."""
    try:
        init_huggingface()
        model = setup_model() 
        interface = setup_gradio_interface(model)
        interface.launch(debug=True)
    except Exception as e:
        print(f"Error during initialization: {str(e)}")
        raise

if __name__ == "__main__":
    main()