medimageinsightmodel.py

"""Medical Image Classification model wrapper class that loads the model, preprocesses inputs and performs inference."""

import torch
from PIL import Image
import pandas as pd
from typing import List, Tuple
import os
import tempfile
import base64
import io

from MedImageInsight.UniCLModel import build_unicl_model
from MedImageInsight.Utils.Arguments import load_opt_from_config_files
from MedImageInsight.ImageDataLoader import build_transforms
from MedImageInsight.LangEncoder import build_tokenizer


class MedImageInsight:
    """Wrapper class for medical image classification model."""

    def __init__(
            self,
            model_dir: str,
            vision_model_name: str,
            language_model_name: str
    ) -> None:
        """Initialize the medical image classifier.

        Args:
            model_dir: Directory containing model files and config
            vision_model_name: Name of the vision model
            language_model_name: Name of the language model
        """
        self.model_dir = model_dir
        self.vision_model_name = vision_model_name
        self.language_model_name = language_model_name
        self.model = None
        self.device = None
        self.tokenize = None
        self.preprocess = None
        self.opt = None

    def load_model(self) -> None:
        """Load the model and necessary components."""
        try:
            # Load configuration
            config_path = os.path.join(self.model_dir, 'config.yaml')
            self.opt = load_opt_from_config_files([config_path])

            # Set paths
            self.opt['LANG_ENCODER']['PRETRAINED_TOKENIZER'] = os.path.join(
                self.model_dir,
                'language_model',
                'clip_tokenizer_4.16.2'
            )
            self.opt['UNICL_MODEL']['PRETRAINED'] = os.path.join(
                self.model_dir,
                'vision_model',
                self.vision_model_name
            )

            # Initialize components
            self.preprocess = build_transforms(self.opt, False)
            self.model = build_unicl_model(self.opt)

            # Set device
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            self.model.to(self.device)

            # Load tokenizer
            self.tokenize = build_tokenizer(self.opt['LANG_ENCODER'])
            self.max_length = self.opt['LANG_ENCODER']['CONTEXT_LENGTH']

            print(f"Model loaded successfully on device: {self.device}")

        except Exception as e:
            print("Failed to load the model:")
            raise e

    @staticmethod
    def decode_base64_image(base64_str: str) -> Image.Image:
        """Decode base64 string to PIL Image and ensure RGB format.

        Args:
            base64_str: Base64 encoded image string

        Returns:
            PIL Image object in RGB format
        """
        try:
            # Remove header if present
            if ',' in base64_str:
                base64_str = base64_str.split(',')[1]

            image_bytes = base64.b64decode(base64_str)
            image = Image.open(io.BytesIO(image_bytes))

            # Convert grayscale (L) or grayscale with alpha (LA) to RGB
            if image.mode in ('L', 'LA'):
                image = image.convert('RGB')

            return image
        except Exception as e:
            raise ValueError(f"Failed to decode base64 image: {str(e)}")

    def predict(self, images: List[str], labels: List[str], multilabel: bool = False) -> List[dict]:
        """Perform zero shot classification on the input images.

        Args:
            images: List of base64 encoded image strings
            labels: List of candidate labels for classification

        Returns:
            DataFrame with columns ["probabilities", "labels"]
        """
        if not self.model:
            raise RuntimeError("Model not loaded. Call load_model() first.")

        if not labels:
            raise ValueError("No labels provided")

        # Create temporary directory for processing
        with tempfile.TemporaryDirectory() as tmp_dir:
            # Process images
            image_list = []
            for img_base64 in images:
                try:
                    img = self.decode_base64_image(img_base64)
                    image_list.append(img)
                except Exception as e:
                    raise ValueError(f"Failed to process image: {str(e)}")

            # Run inference
            probs = self.run_inference_batch(image_list, labels, multilabel)
            probs_np = probs.cpu().numpy()
            results = []
            for prob_row in probs_np:
                # Create label-prob pairs and sort by probability
                label_probs = [(label, float(prob)) for label, prob in zip(labels, prob_row)]
                label_probs.sort(key=lambda x: x[1], reverse=True)

                # Create ordered dictionary from sorted pairs
                results.append({
                    label: prob
                    for label, prob in label_probs
                })

            return results

    def encode(self, images: List[str] = None, texts: List[str] = None):

        output = {
            "image_embeddings"  : None,
            "text_embeddings" : None,
        }

        if not self.model:
            raise RuntimeError("Model not loaded. Call load_model() first.")

        if not images and not texts:
            raise  ValueError("You must provide either images or texts")

        if images is not None:
            with tempfile.TemporaryDirectory() as tmp_dir:
                # Process images
                image_list = []
                for img_base64 in images:
                    try:
                        img = self.decode_base64_image(img_base64)
                        image_list.append(img)
                    except Exception as e:
                        raise ValueError(f"Failed to process image: {str(e)}")
            images = torch.stack([self.preprocess(img) for img in image_list]).to(self.device)
            with torch.no_grad():
                output["image_embeddings"] = self.model.encode_image(images).cpu().numpy()

        if texts is not None:
            text_tokens = self.tokenize(
                texts,
                padding='max_length',
                max_length=self.max_length,
                truncation=True,
                return_tensors='pt'
            )

            # Move text tensors to the correct device
            text_tokens = {k: v.to(self.device) for k, v in text_tokens.items()}
            output["text_embeddings"] = self.model.encode_text(text_tokens).cpu().numpy()


        return output

    def run_inference_batch(
            self,
            images: List[Image.Image],
            texts: List[str],
            multilabel: bool = False
    ) -> torch.Tensor:
        """Perform inference on batch of input images.

        Args:
            images: List of PIL Image objects
            texts: List of text labels
            multilabel: If True, use sigmoid for multilabel classification.
                       If False, use softmax for single-label classification.

        Returns:
            Tensor of prediction probabilities
        """
        # Prepare inputs
        images = torch.stack([self.preprocess(img) for img in images]).to(self.device)

        # Process text
        text_tokens = self.tokenize(
            texts,
            padding='max_length',
            max_length=self.max_length,
            truncation=True,
            return_tensors='pt'
        )

        # Move text tensors to the correct device
        text_tokens = {k: v.to(self.device) for k, v in text_tokens.items()}

        # Run inference
        with torch.no_grad():
            outputs = self.model(image=images, text=text_tokens)
            logits_per_image = outputs[0] @ outputs[1].t() * outputs[2]

            if multilabel:
                # Use sigmoid for independent probabilities per label
                probs = torch.sigmoid(logits_per_image)
            else:
                # Use softmax for single-label classification
                probs = logits_per_image.softmax(dim=1)

        return probs