src/model_loader.py

from pathlib import Path

from transformers import AutoImageProcessor, AutoModelForImageClassification


BASE_DIR = Path(__file__).resolve().parents[1]
DEFAULT_MODEL_DIR = BASE_DIR / "models" / "aishrica_food_predictor"


def load_food_classifier(model_dir: Path = DEFAULT_MODEL_DIR):
    """Load the local ViT image processor and classifier."""
    model_dir = Path(model_dir)
    if not model_dir.exists():
        raise FileNotFoundError(f"Model directory not found: {model_dir}")

    processor = AutoImageProcessor.from_pretrained(
        str(model_dir),
        local_files_only=True,
    )
    model = AutoModelForImageClassification.from_pretrained(
        str(model_dir),
        local_files_only=True,
    )
    model.eval()
    return processor, model