model.py

from pathlib import Path
import io
from typing import Any, Dict, List

from PIL import Image
import torch
from torchvision import models, transforms

# Define path configurations
MODEL_PATH = Path(__file__).resolve().parent / "model_resnet.pth"

# List of all food class names supported by the model
CLASS_NAMES = [
    "apple_pie", "baby_back_ribs", "baklava", "beef_carpaccio", "beef_tartare",
    "beet_salad", "beignets", "bibimbap", "bread_pudding", "breakfast_burrito",
    "bruschetta", "caesar_salad", "cannoli", "caprese_salad", "carrot_cake",
    "ceviche", "cheesecake", "cheese_plate", "chicken_curry", "chicken_quesadilla",
    "chicken_wings", "chocolate_cake", "chocolate_mousse", "churros", "clam_chowder",
    "club_sandwich", "crab_cakes", "creme_brulee", "croque_madame", "cup_cakes"
]

# Mapping classes to human-readable labels and emojis
CLASS_DISPLAY_NAMES = {
    "apple_pie": {"label": "Apple Pie", "icon": "🥧"},
    "baby_back_ribs": {"label": "Baby Back Ribs", "icon": "🍖"},
    "baklava": {"label": "Baklava", "icon": "🍯"},
    "beef_carpaccio": {"label": "Beef Carpaccio", "icon": "🥩"},
    "beef_tartare": {"label": "Beef Tartare", "icon": "🥩"},
    "beet_salad": {"label": "Beet Salad", "icon": "🥗"},
    "beignets": {"label": "Beignets", "icon": "🍩"},
    "bibimbap": {"label": "Bibimbap", "icon": "🍚"},
    "bread_pudding": {"label": "Bread Pudding", "icon": "🍮"},
    "breakfast_burrito": {"label": "Breakfast Burrito", "icon": "🌯"},
    "bruschetta": {"label": "Bruschetta", "icon": "🍅"},
    "caesar_salad": {"label": "Caesar Salad", "icon": "🥗"},
    "cannoli": {"label": "Cannoli", "icon": "🥐"},
    "caprese_salad": {"label": "Caprese Salad", "icon": "🍅"},
    "carrot_cake": {"label": "Carrot Cake", "icon": "🎂"},
    "ceviche": {"label": "Ceviche", "icon": "🍤"},
    "cheesecake": {"label": "Cheesecake", "icon": "🍰"},
    "cheese_plate": {"label": "Cheese Plate", "icon": "🧀"},
    "chicken_curry": {"label": "Chicken Curry", "icon": "🍛"},
    "chicken_quesadilla": {"label": "Chicken Quesadilla", "icon": "🌯"},
    "chicken_wings": {"label": "Chicken Wings", "icon": "🍗"},
    "chocolate_cake": {"label": "Chocolate Cake", "icon": "🍫"},
    "chocolate_mousse": {"label": "Chocolate Mousse", "icon": "🍮"},
    "churros": {"label": "Churros", "icon": "🧋"},
    "clam_chowder": {"label": "Clam Chowder", "icon": "🥣"},
    "club_sandwich": {"label": "Club Sandwich", "icon": "🥪"},
    "crab_cakes": {"label": "Crab Cakes", "icon": "🦀"},
    "creme_brulee": {"label": "Crème Brûlée", "icon": "🍮"},
    "croque_madame": {"label": "Croque Madame", "icon": "🥪"},
    "cup_cakes": {"label": "Cupcakes", "icon": "🧁"},
}
NUM_CLASSES = len(CLASS_NAMES)


# Determine hardware device (GPU or CPU)
def get_device() -> torch.device:
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


# Define standard transformations for input images
def get_transform() -> transforms.Compose:
    return transforms.Compose([
        transforms.Resize((512, 512)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])


# Convert uploaded image bytes to PIL format
def load_image(image_bytes: bytes) -> Image.Image:
    return Image.open(io.BytesIO(image_bytes)).convert("RGB")


# Load the model architecture and pretrained weights
def build_model(device: torch.device) -> torch.nn.Module:
    if not MODEL_PATH.exists():
        raise FileNotFoundError(f"Model file not found: {MODEL_PATH}")

    model = models.resnet50(weights=None)
    in_features = model.fc.in_features
    # Customize the final fully connected layer
    model.fc = torch.nn.Sequential(
        torch.nn.Dropout(p=0.3),
        torch.nn.Linear(in_features=in_features, out_features=NUM_CLASSES)
    )

    state_dict = torch.load(MODEL_PATH, map_location=device)
    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()

    return model


# Perform inference on the provided image
def predict(image: Image.Image, model: torch.nn.Module, device: torch.device, top_k: int = 5) -> List[Dict[str, Any]]:
    transform = get_transform()
    tensor = transform(image).unsqueeze(0).to(device)

    # Disable gradient calculation for inference
    with torch.no_grad():
        outputs = model(tensor)
        probabilities = torch.softmax(outputs[0], dim=0)

    # Get the top K highest probability classes
    top_probs, top_idxs = torch.topk(probabilities, k=min(top_k, NUM_CLASSES))
    predictions = []

    for score, idx in zip(top_probs.tolist(), top_idxs.tolist()):
        class_name = CLASS_NAMES[idx]
        display_info = CLASS_DISPLAY_NAMES.get(class_name, {"label": class_name.replace("_", " ").title(), "icon": "🍽️"})
        predictions.append({
            "label": display_info["label"],
            "icon": display_info["icon"],
            "confidence": float(score * 100),
            "index": int(idx)
        })

    return predictions