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

# Imports locaux
from modeling.BaseModel import BaseModel
from modeling import build_model
from utilities.distributed import init_distributed
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, get_instances

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

def apply_distributed(opt):
    """Applique les paramètres distribués pour le mode multi-processus."""
    print(f"Configuration distribuée appliquée : {opt}")

def init_distributed(opt):
    """Initialize distributed mode without premature CUDA initialization."""
    opt['CUDA'] = opt.get('CUDA', True) and torch.cuda.is_available()
    if 'OMPI_COMM_WORLD_SIZE' not in os.environ:
        # Application started without MPI
        opt['env_info'] = 'no MPI'
        opt['world_size'] = 1
        opt['local_size'] = 1
        opt['rank'] = 0
        opt['local_rank'] = 0  # Ensure this is set to 0
        opt['master_address'] = '127.0.0.1'
        opt['master_port'] = '8673'
    else:
        # Application started with MPI
        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'])

    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:
        opt['device'] = torch.device("cuda", opt['local_rank'])  # Ensure local_rank is integer

    apply_distributed(opt)
    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'] = 'cpu'
    
    pretrained_path = init_huggingface()
    model = BaseModel(opt, build_model(opt))
    state_dict = torch.load(pretrained_path, map_location='cpu', weights_only=True)
    model.load_state_dict(state_dict, strict=False)
    
    # Initialize train_class_names
    model.train_class_names = BIOMED_CLASSES + ["background"]
    
    return model.eval()

import numpy as np
from PIL import Image

def preprocess_image(image):
    """Preprocess image for SEEM model input."""
    if isinstance(image, Image.Image):
        # Convert PIL Image to numpy array
        image = np.array(image)
    
    # Ensure image is float32 and normalized
    image = image.astype(np.float32) / 255.0
    
    # Ensure correct dimensions (B, C, H, W)
    if len(image.shape) == 3:
        image = np.transpose(image, (2, 0, 1))  # HWC -> CHW
        image = np.expand_dims(image, axis=0)  # Add batch dimension
    
    return image

@spaces.GPU 
def predict_image(model, image, prompts):
    """Process image prediction with proper formatting."""
    try:
        # Convert PIL Image to numpy array if needed
        if isinstance(image, Image.Image):
            image = np.array(image)
            
        # Ensure image is in float32 and normalized
        image = image.astype(np.float32) / 255.0
        
        # Transpose from HWC to CHW format
        if len(image.shape) == 3:
            image = np.transpose(image, (2, 0, 1))
            
        # Add batch dimension if needed
        if len(image.shape) == 3:
            image = np.expand_dims(image, axis=0)
            
        # Convert to tensor
        image_tensor = torch.from_numpy(image)
        
        # Move to GPU if available
        if torch.cuda.is_available():
            device = torch.device("cuda", 0)
            model = model.to(device)
            image_tensor = image_tensor.to(device)
        else:
            device = torch.device("cpu")
            
        # Create batched input
        batched_inputs = [{
            "image": image_tensor,
            "prompt": prompts,
            "height": image_tensor.shape[-2],
            "width": image_tensor.shape[-1]
        }]
        
        with torch.no_grad():
            pred_masks = model(batched_inputs)
            
        # Move back to CPU if needed
        if device.type == "cuda":
            model = model.to("cpu")
            pred_masks = [mask.cpu() for mask in pred_masks]
                
        return pred_masks
            
    except Exception as e:
        print(f"Error processing image: {str(e)}")
        raise

def process_image(image, text, model):
    """Process image with proper error handling."""
    try:
        prompts = [p.strip() for p in text.split(',') if p.strip()]
        if not prompts:
            raise ValueError("No valid prompts provided")
            
        pred_masks = predict_image(model, image, prompts)
        
        # Create visualization
        fig = plt.figure(figsize=(5 * (len(pred_masks) + 1), 5))
        
        # Show original image
        plt.subplot(1, len(pred_masks) + 1, 1)
        plt.imshow(preprocess_image(image))
        plt.title("Original")
        plt.axis('off')
        
        # Show predictions
        for i, mask in enumerate(pred_masks):
            plt.subplot(1, len(pred_masks) + 1, i+2)
            plt.imshow(preprocess_image(image))
            plt.imshow(mask.cpu().numpy(), 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(
        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 et spécifiez les éléments à segmenter",
        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()  # Load on CPU
        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()