src/inference/lr_model.py

import json
import joblib
import numpy as np
from PIL import Image


class LRModel:
    """
    Inference pipeline for Logistic Regression model
    trained on 64x64 grayscale flattened images.
    """

    def __init__(self, model_path: str, labels_path: str, image_size: int = 64):
        self.model = joblib.load(model_path)
        self.labels = self._load_labels(labels_path)
        self.image_size = image_size

    def _load_labels(self, labels_path):
        with open(labels_path, "r") as f:
            label_dict = json.load(f)

        # Ensure keys are integer indices, not strings
        label_dict = {int(k): v for k, v in label_dict.items()}
        return label_dict

    def preprocess(self, image: Image.Image) -> np.ndarray:
        """
        Preprocessing matching training:
        - Resize to 64x64
        - Grayscale
        - Normalize to [0,1]
        - Flatten to (1, D)
        """
        img = image.resize((self.image_size, self.image_size))
        img = img.convert("L")  # grayscale
        arr = np.array(img, dtype=np.float32) / 255.0
        arr = arr.reshape(1, -1)  # shape: (1, D)
        return arr

    def predict(self, image: Image.Image):
        """
        Returns:
        {
          "class_id": int,
          "class_name": str,
          "probabilities": {class_name: prob, ...}
        }
        """
        x = self.preprocess(image)
        probs = self.model.predict_proba(x)[0]
        class_id = int(np.argmax(probs))
        class_name = self.labels[class_id]

        # Build probability dict (optional)
        prob_dict = {
            self.labels[i]: float(probs[i]) for i in range(len(probs))
        }

        return {
            "class_id": class_id,
            "class_name": class_name,
            "probabilities": prob_dict
        }