File size: 765 Bytes
50ba6d4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
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