feat: Add files for inference

.dockerignore

@@ -0,0 +1,22 @@
+# Ignore python virtual environments and dependencies
+pytorch-cuda/
+venv/
+.venv/
+env/
+
+# Ignore python cache files
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+
+# Ignore local environment file (production secrets should be injected via env variables)
+.env
+
+# Ignore git folder
+.git/
+.gitignore
+
+# Ignore docker files
+Dockerfile
+.dockerignore

.env.example

@@ -0,0 +1,7 @@
+# Server configuration
+HOST=0.0.0.0
+PORT=8000
+DEBUG_MODE=True
+
+# Hugging Face Token
+HF_TOKEN=your_huggingface_access_token_here

Dockerfile

@@ -0,0 +1,52 @@
+# Use a stable, official Python base image
+FROM python:3.10-slim
+
+# Set environment variables
+# PYTHONUNBUFFERED=1 ensures console logs are printed immediately
+# PYTHONDONTWRITEBYTECODE=1 prevents python from writing .pyc files
+# PORT=7860 is the default port for Hugging Face Spaces
+# HOME=/home/user sets the home folder for the non-root user
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PORT=7860 \
+    HOST=0.0.0.0 \
+    HOME=/home/user
+
+# Install system dependencies required by OpenCV, MediaPipe, and other libraries
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    libgl1 \
+    libglib2.0-0 \
+    libgomp1 \
+    sed \
+    && rm -rf /var/lib/apt/lists/*
+
+# Create a non-root user with UID 1000 (Hugging Face Spaces runs as UID 1000)
+RUN useradd -m -u 1000 user
+WORKDIR /app
+
+# Copy requirements.txt first for build caching
+COPY --chown=user:user requirements.txt /app/
+
+# Remove the custom local PyTorch wheels from requirements.txt to install standard stable versions
+# from PyPI, supporting both CPU and GPU workloads automatically.
+RUN sed -i '/torch==/d' requirements.txt && \
+    sed -i '/torchvision==/d' requirements.txt && \
+    pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir torch torchvision && \
+    pip install --no-cache-dir -r requirements.txt
+
+# Copy the rest of the application files
+COPY --chown=user:user . /app/
+
+# Setup a writeable Hugging Face cache directory inside the home folder of user 1000
+RUN mkdir -p /home/user/.cache/huggingface && chown -R user:user /home/user
+
+# Switch to the non-root user
+USER user
+
+# Expose the default port (Hugging Face Spaces automatically forwards traffic to 7860)
+EXPOSE 7860
+
+# Run the FastAPI server
+CMD ["python", "main.py"]

api/endpoints/predict.py

@@ -0,0 +1,18 @@
+from fastapi import APIRouter
+from schemas.predict import InferenceRequest, PredictionResponse
+from schemas.system import ModelsListResponse
+from services.model_manager import model_manager
+
+router = APIRouter(tags=["Inference"])
+
+@router.get("/models", response_model=ModelsListResponse)
+async def list_models():
+    """List all available models loaded in memory for inference."""
+    return {
+        "available_models": list(model_manager.model_configs.values())
+    }
+
+@router.post("/predict", response_model=PredictionResponse)
+async def predict_personality(request: InferenceRequest):
+    """Predict Big Five personality traits from a base64 encoded face image."""
+    return model_manager.predict(request.model_type, request.image_base64)

api/endpoints/system.py

@@ -0,0 +1,25 @@
+import json
+from fastapi import APIRouter
+from schemas.system import MetadataResponse, HealthResponse
+from services.model_manager import model_manager, DEVICE
+
+router = APIRouter(tags=["System"])
+
+@router.get("/", response_model=MetadataResponse)
+async def root():
+    """Standard root endpoint providing API metadata."""
+    with open("config/metadata.json", "r") as f:
+        metadata = json.load(f)
+    metadata["documentation"] = "/docs"
+    return metadata
+
+@router.get("/health", response_model=HealthResponse)
+async def health_check():
+    """API Health check"""
+    return {
+        "status": "healthy",
+        "device": DEVICE,
+        "models_loaded": list(model_manager.models.keys()),
+        "port": "auto"
+    }
+

api/router.py

@@ -0,0 +1,7 @@
+from fastapi import APIRouter
+from api.endpoints import system, predict
+
+api_router = APIRouter()
+
+api_router.include_router(system.router)
+api_router.include_router(predict.router)

assets/blaze_face_short_range.tflite

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b4578f35940bf5a1a655214a1cce5cab13eba73c1297cd78e1a04c2380b0152f
+size 229746

config/metadata.json

@@ -0,0 +1,6 @@
+{
+  "api_name": "Big Five Personality Inference API",
+  "description": "API for predicting Big Five personality traits (OCEAN) from facial images using deep learning vision models",
+  "version": "1.2.0",
+  "status": "online"
+}

config/models.json

@@ -0,0 +1,26 @@
+[
+    {
+        "id": "vit-b16-augreg-in21k",
+        "name": "ViT-B/16 AugReg (IN21k) - Arch Tuning + Augmentation",
+        "description": "ViT Base patch 16 ArchTuning + AugReg model Trained in 21k images",
+        "repo_id": "lyfesan/vit_base_patch16_224_augreg_in21k_Run_D_The_Ultimate_bigfive"
+    },
+    {
+        "id": "vit-b16-augreg-in21k-ft1k",
+        "name": "ViT-B/16 AugReg (IN21k+1k) - Arch Tuning + Augmentation",
+        "description": "ViT Base patch 16 ArchTuning + AugReg model Trained in 21k images finetuned in 1k images",
+        "repo_id": "lyfesan/vit_base_patch16_224_augreg_in21k_ft_in1k_Run_D_The_Ultimate_bigfive"
+    },
+    {
+        "id": "swinv2-w12-16-archtuning-in22k-ft1k",
+        "name": "SwinV2-B w12-16 (IN22k+1k) - Arch Tuning",
+        "description": "SwinV2 window12-16 ArchTuning model Trained in 22k images finetuned in 1k images",
+        "repo_id": "lyfesan/swinv2_base_window12to16_192to256_ms_in22k_ft_in1k_Run_B_Arch_Tuning_bigfive"
+    },
+    {
+        "id": "swinv2-w16-archtuning-in1k",
+        "name": "SwinV2-B w16 (IN1k) - Arch Tuning",
+        "description": "SwinV2 window16 ArchTuning model trained in 1k images",
+        "repo_id": "lyfesan/swinv2_base_window16_256_ms_in1k_Run_B_Arch_Tuning_bigfive"
+    }
+]

core/config.py

@@ -0,0 +1,16 @@
+import os
+from huggingface_hub import login
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv(override=True)
+
+HOST = os.getenv("HOST", "0.0.0.0")
+PORT = int(os.getenv("PORT", 8000))
+DEBUG_MODE = os.getenv("DEBUG_MODE", "False").lower() in ("true", "1", "t")
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+# Authenticate with Hugging Face
+if HF_TOKEN and HF_TOKEN != "your_huggingface_access_token_here":
+    print("Logging into Hugging Face Hub...")
+    login(token=HF_TOKEN)

core/exceptions.py

@@ -0,0 +1,16 @@
+from fastapi import Request
+from fastapi.responses import JSONResponse
+from starlette.exceptions import HTTPException as StarletteHTTPException
+
+async def custom_404_handler(request: Request, exc: StarletteHTTPException):
+    """Custom handler to format 404 errors cleanly."""
+    if exc.status_code == 404:
+        return JSONResponse(
+            status_code=404,
+            content={
+                "error": "Endpoint not found", 
+                "path": request.url.path,
+                "message": "Please check the URL or visit /docs for available endpoints."
+            }
+        )
+    return JSONResponse(status_code=exc.status_code, content={"error": exc.detail})

main.py

@@ -0,0 +1,56 @@
+from contextlib import asynccontextmanager
+import json
+from fastapi import FastAPI
+from starlette.exceptions import HTTPException as StarletteHTTPException
+
+from core.config import HOST, PORT, DEBUG_MODE
+from core.exceptions import custom_404_handler
+from api.router import api_router
+from services.model_manager import model_manager, DEVICE
+
+# Load metadata from config file
+with open("config/metadata.json", "r") as f:
+    METADATA = json.load(f)
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    print("Downloading/Loading models into VRAM (this takes a moment on first run)...")
+    
+    try:
+        with open("config/models.json", "r") as f:
+            models_config = json.load(f)
+            
+        for model_info in models_config:
+            model_manager.load_hf_model_pipeline(
+                model_info["id"], 
+                model_info["repo_id"], 
+                model_info=model_info
+            )
+    except FileNotFoundError:
+        print("⚠️ models.json not found in config/. No models loaded automatically.")
+
+    yield
+    print("Shutting down API and releasing resources...")
+    model_manager.models.clear()
+    model_manager.transforms_dict.clear()
+    if hasattr(model_manager, "model_configs"):
+        model_manager.model_configs.clear()
+
+app = FastAPI(
+    title=METADATA["api_name"],
+    description=METADATA["description"],
+    version=METADATA["version"],
+    debug=DEBUG_MODE,
+    lifespan=lifespan,
+)
+print(f"API Engine initialized on: {DEVICE.upper()}")
+
+# Register exception handlers
+app.add_exception_handler(StarletteHTTPException, custom_404_handler)
+
+# Include routers
+app.include_router(api_router)
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host=HOST, port=PORT)

requirements.txt

@@ -0,0 +1,76 @@
+aiohappyeyeballs==2.6.1
+aiohttp==3.13.5
+aiosignal==1.4.0
+annotated-doc==0.0.4
+annotated-types==0.7.0
+anyio==4.13.0
+attrs==26.1.0
+certifi==2026.4.22
+charset-normalizer==3.4.7
+click==8.3.3
+colorama==0.4.6
+contourpy==1.3.3
+cycler==0.12.1
+datasets==4.8.5
+dill==0.4.1
+fastapi==0.136.1
+filelock==3.25.2
+fonttools==4.62.1
+frozenlist==1.8.0
+fsspec==2026.2.0
+h11==0.16.0
+hf-xet==1.5.0
+httpcore==1.0.9
+httpx==0.28.1
+huggingface_hub==1.14.0
+idna==3.13
+Jinja2==3.1.6
+joblib==1.5.3
+kiwisolver==1.5.0
+markdown-it-py==4.1.0
+MarkupSafe==3.0.3
+matplotlib==3.10.9      
+mdurl==0.1.2
+mediapipe==0.10.35
+mpmath==1.3.0
+multidict==6.7.1
+multiprocess==0.70.19
+networkx==3.6.1
+numpy==2.4.3
+packaging==26.2
+pandas==3.0.2
+pillow==12.1.1
+propcache==0.4.1
+pyarrow==24.0.0
+pydantic==2.13.4
+pydantic_core==2.46.4
+Pygments==2.20.0
+pyparsing==3.3.2
+python-dateutil==2.9.0.post0
+python-dotenv==1.2.2
+python-multipart==0.0.27
+PyYAML==6.0.3
+requests==2.33.1
+rich==15.0.0
+safetensors==0.7.0
+scikit-learn==1.8.0
+scipy==1.17.1
+seaborn==0.13.2
+setuptools==70.2.0
+shellingham==1.5.4
+six==1.17.0
+starlette==1.0.0
+sympy==1.14.0
+threadpoolctl==3.6.0
+timm==1.0.26
+torch==2.11.0+cu130
+torchvision==0.26.0+cu130
+tqdm==4.67.3
+typer==0.25.1
+typing-inspection==0.4.2
+typing_extensions==4.15.0
+tzdata==2026.2
+urllib3==2.6.3
+uvicorn==0.46.0
+xxhash==3.7.0
+yarl==1.23.0

schemas/predict.py

@@ -0,0 +1,32 @@
+from pydantic import BaseModel, Field
+from typing import Optional
+
+
+class InferenceRequest(BaseModel):
+    """Request body for Big Five personality trait prediction."""
+    model_type: str = Field(
+        ...,
+        description="The ID of the vision model to use for inference",
+        examples=["swinv2", "vit", "pvtv2"],
+    )
+    image_base64: str = Field(
+        ...,
+        description="Base64-encoded image string (JPEG/PNG). Data URI prefix is optional.",
+        examples=["iVBORw0KGgoAAAANSUhEUg..."],
+    )
+
+
+class OCEANTraits(BaseModel):
+    """Big Five (OCEAN) personality trait scores, each ranging from 0.0 to 1.0."""
+    Openness: float = Field(..., ge=0.0, le=1.0, description="Openness to experience", examples=[0.62])
+    Conscientiousness: float = Field(..., ge=0.0, le=1.0, description="Conscientiousness", examples=[0.63])
+    Extraversion: float = Field(..., ge=0.0, le=1.0, description="Extraversion", examples=[0.54])
+    Agreeableness: float = Field(..., ge=0.0, le=1.0, description="Agreeableness", examples=[0.63])
+    Neuroticism: float = Field(..., ge=0.0, le=1.0, description="Neuroticism", examples=[0.60])
+
+
+class PredictionResponse(BaseModel):
+    """Response containing the model used, predicted OCEAN traits, and the cropped face image."""
+    model_used: str = Field(..., description="The ID of the model that produced the prediction", examples=["swinv2"])
+    predictions: OCEANTraits = Field(..., description="Predicted Big Five personality trait scores")
+    cropped_face_base64: Optional[str] = Field(None, description="Base64 encoded cropped face image, if face extraction was used.", examples=["/9j/4AAQSkZJRgABAQEASABIAAD/4..."])

schemas/system.py

@@ -0,0 +1,31 @@
+from pydantic import BaseModel, Field
+from typing import List
+
+
+class MetadataResponse(BaseModel):
+    """API metadata returned by the root endpoint."""
+    api_name: str = Field(..., description="Name of the API", examples=["Big Five Personality Inference API"])
+    description: str = Field(..., description="Brief description of the API", examples=["API for predicting Big Five personality traits"])
+    version: str = Field(..., description="Current API version", examples=["1.2.0"])
+    status: str = Field(..., description="Current API status", examples=["online"])
+    documentation: str = Field(..., description="Path to interactive API docs", examples=["/docs"])
+
+
+class HealthResponse(BaseModel):
+    """Health check response."""
+    status: str = Field(..., description="Health status of the API", examples=["healthy"])
+    device: str = Field(..., description="Compute device in use", examples=["cuda"])
+    models_loaded: List[str] = Field(..., description="List of model IDs currently loaded in memory", examples=[["swinv2", "vit", "pvtv2"]])
+    port: str = Field(..., description="Port configuration", examples=["auto"])
+
+
+class ModelDetail(BaseModel):
+    """Details of an available inference model."""
+    id: str = Field(..., description="The ID of the model to be used in predictions", examples=["swinv2"])
+    name: str = Field(..., description="Human-readable name of the model", examples=["Swin Transformer V2"])
+    description: str = Field(..., description="Description of the model", examples=["SwinV2 Base model optimized for Big Five personality traits prediction"])
+    repo_id: str = Field(..., description="Hugging Face model repository ID", examples=["lyfesan/swinv2_base_..."])
+
+class ModelsListResponse(BaseModel):
+    """List of available inference models."""
+    available_models: List[ModelDetail] = Field(..., description="List of models available for inference")

services/face_extractor.py

@@ -0,0 +1,74 @@
+import math
+import numpy as np
+from PIL import Image
+import mediapipe as mp
+from mediapipe.tasks import python
+from mediapipe.tasks.python import vision
+
+class FaceExtractor:
+    def __init__(self, model_path: str = "assets/blaze_face_short_range.tflite"):
+        self.model_path = model_path
+        base_options = python.BaseOptions(model_asset_path=self.model_path)
+        options = vision.FaceDetectorOptions(
+            base_options=base_options, 
+            running_mode=vision.RunningMode.IMAGE,
+            min_detection_confidence=0.70
+        )
+        self.detector = vision.FaceDetector.create_from_options(options)
+        self.offset_percentage = 0.30
+
+    def extract_main_face(self, pil_image: Image.Image) -> Image.Image:
+        """
+        Detects faces in the given PIL Image, scores them to find the main face,
+        and returns the cropped main face. Returns None if no face is detected.
+        """
+        # Convert PIL Image to numpy array (RGB)
+        frame = np.array(pil_image)
+        
+        img_h, img_w, _ = frame.shape
+        frame_cx, frame_cy = img_w / 2, img_h / 2
+        
+        # Mediapipe requires the image to be in ImageFormat.SRGB
+        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame)
+        results = self.detector.detect(mp_image)
+        
+        if not results.detections:
+            return None
+            
+        best_face_bbox = None
+        highest_score = -float('inf')
+        
+        for detection in results.detections:
+            bbox = detection.bounding_box
+            confidence = detection.categories[0].score
+            x, y, w, h = bbox.origin_x, bbox.origin_y, bbox.width, bbox.height
+            face_cx, face_cy = x + (w / 2), y + (h / 2)
+            
+            area = w * h
+            distance_to_center = math.sqrt((frame_cx - face_cx)**2 + (frame_cy - face_cy)**2)
+            score = (area * confidence) - (distance_to_center * 50) 
+            
+            if score > highest_score:
+                highest_score = score
+                best_face_bbox = (x, y, w, h)
+                
+        if not best_face_bbox:
+            return None
+            
+        # Crop with offset
+        x, y, w, h = best_face_bbox
+        offset_w = int(w * self.offset_percentage)
+        offset_h = int(h * self.offset_percentage)
+        
+        new_x = max(0, x - offset_w)
+        new_y = max(0, y - offset_h)
+        new_w = min(img_w - new_x, w + (2 * offset_w))
+        new_h = min(img_h - new_y, h + (2 * offset_h))
+        
+        cropped_face_np = frame[new_y:new_y+new_h, new_x:new_x+new_w]
+        
+        # Convert back to PIL Image
+        if cropped_face_np.size > 0:
+            return Image.fromarray(cropped_face_np)
+            
+        return None

services/inference.py

@@ -0,0 +1,27 @@
+import torch.nn as nn
+import timm
+from huggingface_hub import PyTorchModelHubMixin
+
+class BigFiveRegressor(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, timm_name, use_complex_head=True):
+        super().__init__()
+        self.backbone = timm.create_model(timm_name, pretrained=False, num_classes=0) 
+        num_features = self.backbone.num_features
+        
+        if use_complex_head:
+            self.regression_head = nn.Sequential(
+                nn.Linear(num_features, 512),
+                nn.GELU(),
+                nn.Dropout(0.3),
+                nn.Linear(512, 5),
+                nn.Sigmoid()
+            )
+        else:
+            self.regression_head = nn.Sequential(
+                nn.Linear(num_features, 5),
+                nn.Sigmoid()
+            )
+            
+    def forward(self, x):
+        features = self.backbone(x)
+        return self.regression_head(features)

services/model_manager.py

@@ -0,0 +1,104 @@
+import io
+import base64
+import torch
+from PIL import Image
+from torchvision import transforms
+from fastapi import HTTPException
+from services.inference import BigFiveRegressor
+from schemas.predict import OCEANTraits, PredictionResponse
+from services.face_extractor import FaceExtractor
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+# DEVICE = "cpu"
+class ModelManager:
+    def __init__(self):
+        self.models = {}
+        self.transforms_dict = {}
+        self.model_configs = {}
+        try:
+            self.face_extractor = FaceExtractor()
+        except Exception as e:
+            print(f"Warning: Failed to initialize FaceExtractor: {e}")
+            self.face_extractor = None
+
+    def load_hf_model_pipeline(self, model_key: str, repo_id: str, model_info: dict = None):
+        """Loads model from Hugging Face and creates its specific preprocessing transform."""
+        try:
+            model = BigFiveRegressor.from_pretrained(repo_id)
+            model.to(DEVICE)
+            model.eval()
+
+            # SwinV2 uses 256x256, ViT/PVTv2 use 224x224
+            IMG_SIZE = 256 if 'swinv2' in model_key else 224
+            transform = transforms.Compose([
+                transforms.Resize((IMG_SIZE, IMG_SIZE)),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+            ])
+            
+            self.models[model_key] = model
+            self.transforms_dict[model_key] = transform
+            if model_info:
+                self.model_configs[model_key] = model_info
+            print(f"✅ Loaded {model_key.upper()} from {repo_id}")
+        except Exception as e:
+            print(f"⚠️ Failed to load {model_key} from {repo_id}. Error: {e}")
+
+    def predict(self, model_type: str, image_base64: str) -> PredictionResponse:
+        model_type_lower = model_type.lower()
+        if model_type_lower not in self.models:
+            raise HTTPException(status_code=400, detail=f"Invalid model type. Choose from: {list(self.models.keys())}")
+
+        # Decode Base64 to Image
+        try:
+            # Strip header if frontend accidentally includes "data:image/jpeg;base64,"
+            base64_data = image_base64.split(",")[-1] 
+            image_data = base64.b64decode(base64_data)
+            image = Image.open(io.BytesIO(image_data)).convert("RGB")
+        except Exception:
+            raise HTTPException(status_code=400, detail="Invalid Base64 image payload.")
+
+        # Face Extraction
+        cropped_base64 = None
+        if self.face_extractor:
+            image = self.face_extractor.extract_main_face(image)
+            if image is None:
+                raise HTTPException(status_code=400, detail="No face detected in the image.")
+            
+            # Convert back to base64 for response
+            buffered = io.BytesIO()
+            image.save(buffered, format="JPEG")
+            cropped_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+        # Transform and Infer
+        transform = self.transforms_dict[model_type_lower]
+        input_tensor = transform(image).unsqueeze(0).to(DEVICE)
+        
+        model = self.models[model_type_lower]
+        with torch.no_grad():
+            with torch.amp.autocast('cuda' if DEVICE == 'cuda' else 'cpu'):
+                output = model(input_tensor)
+                probabilities = output.squeeze().cpu().to(torch.float32).numpy()
+
+        # 1. Map the raw array to the order the model was trained on
+        raw_traits = ['Extraversion', 'Neuroticism', 'Agreeableness', 'Conscientiousness', 'Openness']
+        raw_results = {trait: float(score) for trait, score in zip(raw_traits, probabilities)}
+
+        # 2. Standardize to the OCEAN format using Pydantic
+        standardized_ocean = OCEANTraits(
+            Openness=raw_results['Openness'],
+            Conscientiousness=raw_results['Conscientiousness'],
+            Extraversion=raw_results['Extraversion'],
+            Agreeableness=raw_results['Agreeableness'],
+            Neuroticism=raw_results['Neuroticism']
+        )
+
+        # 3. Return the strictly formatted Pydantic object
+        return PredictionResponse(
+            model_used=model_type_lower, 
+            predictions=standardized_ocean,
+            cropped_face_base64=cropped_base64
+        )
+
+# Global instance to be used across the application
+model_manager = ModelManager()