Push Files

.gitignore

@@ -0,0 +1,3 @@
+venv/
+.env
+__pycache__/

__init__.py

@@ -0,0 +1,2 @@
+import sys
+sys.path.append("..")

embedding_generator.py

@@ -0,0 +1,86 @@
+from huggingface_hub import login, from_pretrained_keras
+
+import os
+import glob
+import time
+import h5py
+import numpy as np
+# import pandas as pd
+from PIL import Image
+from tqdm import tqdm
+import tensorflow as tf
+from dotenv import load_dotenv
+
+load_dotenv()
+
+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(token=hf_token)
+
+def load_model():
+    """Load PathFoundation model from Hugging Face"""
+    print("Loading PathFoundation model...")
+    model = from_pretrained_keras("google/path-foundation")
+    infer = model.signatures["serving_default"]
+    print("Model loaded!")
+    return infer
+
+def load_model():
+    """Load PathFoundation model from Hugging Face"""
+    print("Loading PathFoundation model...")
+    import tensorflow as tf
+    import keras
+    from huggingface_hub import snapshot_download
+    
+    # Download the model from HuggingFace
+    model_path = snapshot_download(repo_id="google/path-foundation")
+    
+    # Load as TFSMLayer
+    model = keras.layers.TFSMLayer(
+        model_path,
+        call_endpoint='serving_default'
+    )
+    print("Model loaded!")
+    return model
+
+
+def process_image(image_input, infer_function):
+    """Process a single image and get embedding
+    
+    Args:
+        image_input: Either a file path (str) or image data (bytes/BytesIO/numpy array)
+        infer_function: The model inference function
+    
+    Returns:
+        Embedding vector or None if processing fails
+    """
+    try:
+        # Handle different input types
+        if isinstance(image_input, str):
+            # It's a file path
+            img = Image.open(image_input).convert('RGB')
+        elif isinstance(image_input, bytes) or hasattr(image_input, 'read'):
+            # It's image data from frontend (bytes or BytesIO)
+            img = Image.open(image_input).convert('RGB')
+        elif isinstance(image_input, np.ndarray):
+            # It's already a numpy array
+            img = Image.fromarray(image_input.astype('uint8')).convert('RGB')
+        else:
+            raise ValueError(f"Unsupported image input type: {type(image_input)}")
+        
+        # Resize to 224x224 if needed
+        if img.size != (224, 224):
+            img = img.resize((224, 224))
+        # Convert to tensor and normalize
+        tensor = tf.cast(tf.expand_dims(np.array(img), axis=0), tf.float32) / 255.0
+        
+        # Get embedding
+        embeddings = infer_function(tf.constant(tensor))
+
+        embedding_vector = embeddings['output_0'].numpy().flatten()
+   
+        return embedding_vector
+    except Exception as e:
+        print(f"Error processing image: {e}")
+        return None

main.py

@@ -0,0 +1,49 @@
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi.responses import JSONResponse
+import uvicorn
+from typing import List
+import os
+import numpy as np
+from PIL import Image
+import io
+from embedding_generator import load_model, process_image
+
+app = FastAPI(title="Medical Image Embedding Generator")
+
+
+global infer
+infer = load_model()
+
+
+@app.post("/embeddings")
+async def generate_embeddings(file: UploadFile = File(...)):
+    """
+    Upload a medical image (JPEG, PNG, TIFF) and get embeddings
+    """
+    content_type = file.content_type
+    if not (content_type.startswith("image/") or 
+            file.filename.endswith((".tif", ".tiff", ".jpg", ".jpeg", ".png", ".bmp"))):
+        raise HTTPException(status_code=400, detail="File must be an image (JPEG, PNG, BMP) or TIFF format")
+    
+    try:
+        # Read the file content
+        embedding = process_image(file.file, infer)
+        if embedding is None:
+            raise HTTPException(status_code=500, detail="Error processing image")
+        
+        return_content = {
+            "filename": file.filename,
+            "embedding": embedding.tolist(),
+        }
+        
+        return JSONResponse(content=return_content)
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
+
+@app.get("/")
+async def root():
+    return {"message": "Welcome to Medical Image Embedding Generator API. Use /embeddings endpoint to upload images."}
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)

requirements.txt

@@ -0,0 +1,4 @@
+tf-nightly[and-cuda]
+python-dotenv
+pillow
+huggingface-hub