Add files using upload-large-folder tool

.gitignore

@@ -0,0 +1,10 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv

.python-version

@@ -0,0 +1 @@
+3.8.19

README.md

@@ -1,3 +0,0 @@
----
-license: mit
----

example.py

@@ -0,0 +1,44 @@
+# Initialize classifier
+from medimageinsightmodel import MedImageInsight
+import base64
+
+
+classifier = MedImageInsight(
+    model_dir="2024.09.27",
+    vision_model_name="medimageinsigt-v1.0.0.pt",
+    language_model_name="language_model.pth"
+)
+
+def read_image(image_path):
+    with open(image_path, "rb") as f:
+        return f.read()
+
+# Load model
+classifier.load_model()
+
+import urllib.request
+
+image_url = "https://openi.nlm.nih.gov/imgs/512/145/145/CXR145_IM-0290-1001.png"
+image_path = "CXR145_IM-0290-1001.png"
+
+urllib.request.urlretrieve(image_url, image_path)
+print(f"Image downloaded to {image_path}")
+
+
+image = base64.encodebytes(read_image(image_path)).decode("utf-8")
+
+# Example inference
+images = [image]
+labels = ["normal", "Pneumonia", "unclear"]
+
+#Zero-shot classification
+results = classifier.predict(images, labels)
+print(results)
+
+#Image embeddings
+results = classifier.encode(images = images)
+print(results)
+
+#Text embeddings
+results = classifier.encode(texts = labels)
+print(results)

flask_app.py

@@ -0,0 +1,57 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List
+import uvicorn
+from medimageinsightmodel import MedImageInsight
+import base64
+
+# Initialize FastAPI app
+app = FastAPI(title="Medical Image Analysis API")
+
+# Initialize model
+classifier = MedImageInsight(
+    model_dir="2024.09.27",
+    vision_model_name="medimageinsigt-v1.0.0.pt",
+    language_model_name="language_model.pth"
+)
+classifier.load_model()
+
+
+class ClassificationRequest(BaseModel):
+    images: List[str]  # Base64 encoded images
+    labels: List[str]
+    multilabel : bool = False
+
+class EmbeddingRequest(BaseModel):
+    images: List[str] = None  # Base64 encoded images
+    texts: List[str] = None
+
+@app.post("/predict")
+async def predict(request: ClassificationRequest):
+    try:
+        results = classifier.predict(
+            images=request.images,
+            labels=request.labels,
+            multilabel = request.multilabel
+        )
+        return {"predictions": results}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/encode")
+async def encode(request: EmbeddingRequest):
+    try:
+        results = classifier.encode(images=request.images, texts= request.texts)
+        results["image_embeddings"] = results["image_embeddings"].tolist() if results["image_embeddings"] is not None else None
+        results["text_embeddings"] = results["text_embeddings"].tolist() if results["text_embeddings"] is not None else None
+
+        return results
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/health")
+async def health():
+    return {"status": "healthy"}
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)

medimageinsightmodel.py

@@ -0,0 +1,239 @@
+"""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
+
+

pyproject.toml

@@ -0,0 +1,29 @@
+[project]
+name = "MedImageInsights"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = "==3.8.19"
+dependencies = [
+    "mlflow==2.14.3",
+    "cffi==1.17.1",
+    "cloudpickle==3.0.0",
+    "colorama==0.4.6",
+    "einops==0.8.0",
+    "ftfy==6.2.3",
+    "fvcore==0.1.5.post20221221",
+    "mup==1.0.0",
+    "numpy==1.24.4",
+    "packaging==24.1",
+    "pandas==2.0.3",
+    "pyyaml==6.0.2",
+    "requests==2.32.3",
+    "sentencepiece==0.2.0",
+    "tenacity==9.0.0",
+    "timm==1.0.9",
+    "tornado==6.4.1",
+    "transformers==4.46.0",
+#    "huggingface-hub==0.26.1",
+    "fastapi[standard]>=0.115.3",
+#    "opencv-python>=4.10.0.84",
+]

requirements.txt

@@ -0,0 +1,18 @@
+mlflow==2.14.3
+cffi==1.17.1
+cloudpickle==3.0.0
+colorama==0.4.6
+einops==0.8.0
+ftfy==6.2.3
+fvcore==0.1.5.post20221221
+mup==1.0.0
+numpy==1.24.4
+packaging==24.1
+pandas==2.0.3
+pyyaml==6.0.2
+requests==2.32.3
+sentencepiece==0.2.0
+tenacity==9.0.0
+timm==1.0.9
+tornado==6.4.1
+transformers==4.16.2