feat(server): implement FastAPI application with model loading(HF HUB), CORS support, prediction endpoint and Docker

.dockerignore

@@ -0,0 +1,46 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+.venv/
+venv/
+ENV/
+env/
+
+# Development
+.env
+.env.*
+.git/
+.gitignore
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Logs
+*.log
+
+# Temporary files
+*.tmp
+*.temp
+
+# Test files
+tests/
+test_*
+*_test.py
+
+# Coverage
+.coverage
+htmlcov/
+
+# Documentation
+README.md
+*.md

Dockerfile

@@ -0,0 +1,35 @@
+# apps/server/Dockerfile
+
+FROM python:3.11-slim AS base
+ENV PYTHONDONTWRITEBYTECODE=1 PYTHONUNBUFFERED=1
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    && rm -rf /var/lib/apt/lists/*
+WORKDIR /app
+
+FROM base AS builder
+COPY requirements.txt /app/requirements.txt
+RUN pip install --upgrade pip && \
+    pip wheel --no-cache-dir --wheel-dir /app/wheels -r /app/requirements.txt
+
+FROM python:3.11-slim AS runtime
+ENV PYTHONDONTWRITEBYTECODE=1 PYTHONUNBUFFERED=1
+# add runtime libs you actually need; xgboost often needs libgomp1
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    libgomp1 \
+    && rm -rf /var/lib/apt/lists/*
+WORKDIR /app
+
+# deps
+COPY --from=builder /app/wheels /wheels
+RUN pip install --no-cache-dir /wheels/* && rm -rf /wheels
+
+# app source
+COPY . /app
+
+# non-root
+RUN useradd -m appuser
+USER appuser
+
+EXPOSE 8000
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

app/api/routes.py

@@ -0,0 +1,72 @@
+import io
+import numpy as np
+import torch
+from PIL import Image
+from fastapi import APIRouter, Depends, File, UploadFile, Form, HTTPException
+from fastapi.responses import JSONResponse
+from typing import Optional
+
+from ..state import app_state
+from ..utils import normalize_age
+
+from fastapi_limiter.depends import RateLimiter
+
+from app.config import RATE_TIMES, RATE_SECONDS
+
+router = APIRouter()
+
+@router.get("/health")
+def health():
+    return {
+        "status": "ok", 
+        "device": str(app_state.device), 
+        "classes": app_state.id2label, 
+        "model_loaded": app_state.is_model_loaded()
+    }
+
+
+@router.post("/predict", dependencies=[Depends(RateLimiter(times=RATE_TIMES, seconds=RATE_SECONDS))],)
+async def predict(
+    file: UploadFile = File(..., description="RGB lesion image"),
+    age: Optional[float] = Form(None),
+    localization: Optional[str] = Form("unknown"),
+    top_k: Optional[int] = Form(3),
+):
+    if not app_state.is_model_loaded():
+        raise HTTPException(status_code=503, detail="Model not loaded yet")
+
+    # Read image
+    try:
+        img_bytes = await file.read()
+        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Invalid image: {e}")
+
+    # Preprocess image
+    px = app_state.image_processor(img, return_tensors="pt")["pixel_values"].to(app_state.device)
+
+    # Tabular vector
+    loc = (localization or "unknown").strip().lower()
+    loc_oh = app_state.loc_encoder.transform(np.array([loc]).reshape(-1, 1))  # (1, L)
+    norm_age = normalize_age(age, app_state.age_stats["age_min"], app_state.age_stats["age_max"], app_state.age_stats["age_mean"])
+    tab = np.concatenate([loc_oh, np.array([[norm_age]])], axis=1).astype("float32")
+    tab_t = torch.tensor(tab, dtype=torch.float32, device=app_state.device)
+
+    # Forward
+    with torch.no_grad():
+        logits = app_state.model(pixel_values=px, tabular_features=tab_t)
+        probs = torch.softmax(logits, dim=-1).cpu().numpy()[0]
+
+    # Top-k
+    k = max(1, min(int(top_k or 3), len(probs)))
+    idxs = np.argsort(-probs)[:k]
+    top = [{"label": app_state.id2label[int(i)], "probability": float(probs[i])} for i in idxs]
+    dist = {app_state.id2label[int(i)]: float(p) for i, p in enumerate(probs)}
+
+    payload = {
+        "top": top
+        # "distribution": dist,
+        # "accepted_localizations_example": app_state.valid_localizations[:10]
+    }
+
+    return JSONResponse(content=payload)

app/config.py

@@ -0,0 +1,14 @@
+import os
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
+
+# Redis config
+REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379")
+
+# Rate limiting: requests per time window per IP
+RATE_TIMES = int(os.getenv("RATE_TIMES", "60"))
+RATE_SECONDS = int(os.getenv("RATE_SECONDS", "60"))
+
+TRUSTED_HOSTS = os.getenv("TRUSTED_HOSTS", "*").split(",")

app/models.py

@@ -0,0 +1,57 @@
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, PretrainedConfig, AutoModel
+
+
+class SkinCancerConfig(PretrainedConfig):
+    model_type = "vit_tabular_skin_cancer"
+    
+    def __init__(self,
+                 vision_model_checkpoint="google/vit-base-patch16-224-in21k",
+                 tabular_dim=0,
+                 num_labels=7,
+                 id2label=None, 
+                 label2id=None,
+                 age_min=0.0, 
+                 age_max=100.0, 
+                 age_mean=50.0,
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.vision_model_checkpoint = vision_model_checkpoint
+        self.tabular_dim = tabular_dim
+        self.num_labels = num_labels
+        self.id2label = id2label
+        self.label2id = label2id
+        self.age_min = age_min
+        self.age_max = age_max
+        self.age_mean = age_mean
+
+
+class SkinCancerViT(PreTrainedModel):
+    config_class = SkinCancerConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.vision = AutoModel.from_pretrained(config.vision_model_checkpoint)
+        hdim = self.vision.config.hidden_size
+
+        self.tabular = nn.Sequential(
+            nn.Linear(config.tabular_dim, 128),
+            nn.ReLU(),
+            nn.Dropout(0.1),
+            nn.Linear(128, 64),
+            nn.ReLU()
+        )
+        self.classifier = nn.Linear(hdim + 64, config.num_labels)
+        self.post_init()
+
+    def forward(self, pixel_values, tabular_features):
+        vout = self.vision(pixel_values=pixel_values, output_hidden_states=False, return_dict=True)
+        if getattr(vout, "pooler_output", None) is not None:
+            vfeat = vout.pooler_output
+        else:
+            vfeat = vout.last_hidden_state[:, 0, :]  # CLS
+        tfeat = self.tabular(tabular_features.float())
+        feats = torch.cat([vfeat, tfeat], dim=-1)
+        logits = self.classifier(feats)
+        return logits

app/services.py

@@ -0,0 +1,132 @@
+import os
+import torch
+import numpy as np
+from transformers import AutoImageProcessor, AutoConfig
+from sklearn.preprocessing import OneHotEncoder
+from huggingface_hub import hf_hub_download, list_repo_files
+
+from fastapi import Request
+
+from .state import app_state
+from .models import SkinCancerConfig, SkinCancerViT
+from .utils import load_json
+
+
+def load_model():
+    """Load and initialize the model and related components from Hugging Face."""
+    print(f"Loading model from Hugging Face: {app_state.HF_REPO_ID}")
+    
+    try:
+        # Download and load label maps from HF
+        print("Loading label maps...")
+        label2id_path = hf_hub_download(repo_id=app_state.HF_REPO_ID, filename="label2id.json")
+        id2label_path = hf_hub_download(repo_id=app_state.HF_REPO_ID, filename="id2label.json")
+        
+        app_state.label2id = load_json(label2id_path)
+        id2label_raw = load_json(id2label_path)
+        app_state.id2label.update({int(k): v for k, v in id2label_raw.items()})
+        print(f"Loaded {len(app_state.id2label)} classes")
+
+        # Download and load encoder categories
+        print("Loading encoder categories...")
+        cats_path = hf_hub_download(repo_id=app_state.HF_REPO_ID, filename="loc_encoder_categories.npy")
+        cats = np.load(cats_path, allow_pickle=True)
+        app_state.loc_encoder = OneHotEncoder(handle_unknown="ignore", sparse_output=False)
+        app_state.loc_encoder.fit(np.array(cats).reshape(-1, 1))
+        app_state.valid_localizations[:] = list(cats.tolist())
+        print(f"Loaded {len(app_state.valid_localizations)} localizations")
+
+        # Tabular dim = one-hot length + 1 (age)
+        app_state.tab_dim = app_state.loc_encoder.transform(np.array(["unknown"]).reshape(-1, 1)).shape[1] + 1
+        print(f"Tabular dimension: {app_state.tab_dim}")
+
+        # Download and load age stats
+        age_stats_path = hf_hub_download(repo_id=app_state.HF_REPO_ID, filename="age_stats.json")
+        app_state.age_stats.update(load_json(age_stats_path))
+        print(f"Age stats: {app_state.age_stats}")
+
+        # Download and read the HF config to get the vision backbone name
+        print("Loading model config...")
+        config_path = hf_hub_download(repo_id=app_state.HF_REPO_ID, filename="best_model/config.json")
+        cfg_json = load_json(config_path)
+        app_state.vision_ckpt = cfg_json.get("vision_model_checkpoint", app_state.vision_ckpt)
+        print(f"Vision checkpoint: {app_state.vision_ckpt}")
+        
+        app_state.image_processor = AutoImageProcessor.from_pretrained(app_state.vision_ckpt)
+        print("Image processor loaded")
+
+        # Build model config
+        print("Building model config...")
+        sc_cfg = SkinCancerConfig(
+            vision_model_checkpoint=app_state.vision_ckpt,
+            tabular_dim=app_state.tab_dim,
+            num_labels=len(app_state.id2label),
+            id2label=app_state.id2label,
+            label2id=app_state.label2id,
+            age_min=app_state.age_stats["age_min"],
+            age_max=app_state.age_stats["age_max"],
+            age_mean=app_state.age_stats["age_mean"]
+        )
+        
+        # Initialize empty model with our config
+        print("Initializing model...")
+        model_init = SkinCancerViT(sc_cfg)
+        
+        # Load weights from HF
+        print("Loading model weights from Hugging Face...")
+        try:
+            # Try to load from safetensors first
+            model_path = hf_hub_download(repo_id=app_state.HF_REPO_ID, filename="best_model/model.safetensors")
+            from safetensors.torch import load_file as safe_load
+            print(f"Loading from safetensors: {model_path}")
+            state = safe_load(model_path)
+        except Exception as e:
+            print(f"Safetensors not found, trying pytorch_model.bin: {e}")
+            model_path = hf_hub_download(repo_id=app_state.HF_REPO_ID, filename="best_model/pytorch_model.bin")
+            state = torch.load(model_path, map_location="cpu")
+
+        # Remove training-only keys like loss_fct.weight
+        to_drop = [k for k in list(state.keys()) if k.startswith("loss_fct.")]
+        for k in to_drop:
+            state.pop(k, None)
+
+        # Load with strict=False to ignore harmless mismatches
+        missing, unexpected = model_init.load_state_dict(state, strict=False)
+        if unexpected:
+            print("Ignored unexpected keys:", unexpected)
+        if missing:
+            print("Missing keys:", missing)
+
+        print(f"Using device: {app_state.device}")
+        model_init.to(app_state.device)
+        model_init.eval()
+        app_state.model = model_init
+        print("Model loaded successfully from Hugging Face!")
+
+        # Patch size / grid (if available from vision config)
+        try:
+            app_state.vit_patch_size = getattr(model_init.vision.config, "patch_size", app_state.vit_patch_size)
+            # For square inputs (224×224) with non-overlapping patches
+            size = app_state.image_processor.size
+            if isinstance(size, dict):
+                h = size.get("height", 224)
+                w = size.get("width", 224)
+            else:
+                h = w = size
+            app_state.vit_grid = (h // app_state.vit_patch_size, w // app_state.vit_patch_size)
+            print(f"ViT grid: {app_state.vit_grid}")
+        except Exception as e:
+            print(f"Error setting ViT grid: {e}")
+            app_state.vit_patch_size, app_state.vit_grid = app_state.DEFAULT_VIT_PATCH_SIZE, app_state.DEFAULT_VIT_GRID
+            
+    except Exception as e:
+        print(f"Error loading model from Hugging Face: {e}")
+        raise
+
+
+async def get_client_ip(request: Request) -> str:
+    # First hop of X-Forwarded-For is original client. Fall back to direct socket IP.
+    xff = request.headers.get("x-forwarded-for")
+    if xff:
+        return xff.split(",")[0].strip()
+    return request.client.host

app/state.py

@@ -0,0 +1,50 @@
+import os
+from typing import Dict, List, Optional
+import torch
+from transformers import AutoImageProcessor
+from sklearn.preprocessing import OneHotEncoder
+
+from .models import SkinCancerViT
+
+
+class AppState:
+    """Centralized state management for the application."""
+    
+    def __init__(self):
+        # Hugging Face configuration
+        # Hi there human. What you looking at.
+        self.HF_REPO_ID = os.environ.get("HF_REPO_ID", "HelloWorld47474747/skin_vit_tabular")
+        
+        # Default settings
+        self.DEFAULT_AGE_STATS = {"age_min": 0.0, "age_max": 100.0, "age_mean": 50.0}
+        self.DEFAULT_VIT_PATCH_SIZE = 16
+        self.DEFAULT_VIT_GRID = (14, 14)
+        self.DEFAULT_VISION_CKPT = "google/vit-base-patch16-224-in21k"
+        
+        # Model state
+        self.image_processor: Optional[AutoImageProcessor] = None
+        self.model: Optional[SkinCancerViT] = None
+        self.label2id: Dict[str, int] = {}
+        self.id2label: Dict[int, str] = {}
+        self.loc_encoder: Optional[OneHotEncoder] = None
+        self.age_stats = self.DEFAULT_AGE_STATS.copy()
+        self.tab_dim = 0
+        self.valid_localizations: List[str] = []
+        self.vit_patch_size = self.DEFAULT_VIT_PATCH_SIZE
+        self.vit_grid = self.DEFAULT_VIT_GRID
+        self.vision_ckpt = self.DEFAULT_VISION_CKPT
+        
+        # Device
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    def is_model_loaded(self) -> bool:
+        """Check if the model is loaded."""
+        return self.model is not None and self.image_processor is not None
+    
+    def get_device(self) -> torch.device:
+        """Get the current device."""
+        return self.device
+
+
+# Global state instance
+app_state = AppState()

app/utils.py

@@ -0,0 +1,22 @@
+import json
+import numpy as np
+from typing import Optional
+
+
+def load_json(path: str) -> dict:
+    """Load JSON file."""
+    with open(path, "r") as f:
+        return json.load(f)
+
+
+def normalize_age(age: Optional[float], amin: float, amax: float, amean: float) -> float:
+    """Normalize age to [0, 1] range."""
+    if age is None:
+        age = amean
+    try:
+        age = float(age)
+    except Exception:
+        age = amean
+    if amax == amin:
+        return 0.0
+    return (age - amin) / (amax - amin)

main.py

@@ -1,7 +1,43 @@
+import os
+from contextlib import asynccontextmanager
 from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
 
-app = FastAPI()
+from app.services import load_model, get_client_ip
+from app.api.routes import router
 
-@app.get("/health")
-def read_root():
-    return {"message": "Hello?"}
+from redis import asyncio as redis
+from fastapi_limiter import FastAPILimiter
+from app.config import REDIS_URL, TRUSTED_HOSTS
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Startup
+    print("Loading model...")
+    load_model()
+    print("Model loaded successfully!")
+    redis_client = redis.from_url(REDIS_URL, encoding="utf-8", decode_responses=True)
+    await FastAPILimiter.init(redis_client, identifier=get_client_ip)
+    yield
+    # Shutdown (if needed)
+
+
+app = FastAPI(title="Skin Cancer ViT+Tabular API", lifespan=lifespan)
+
+# CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=TRUSTED_HOSTS if "*" not in TRUSTED_HOSTS else ["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Include API routes
+app.include_router(router, prefix="/api")
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run("main:app", host="0.0.0.0", port=int(os.environ.get("PORT", 8000)))