Spaces:
Sleeping
Sleeping
| import torch | |
| from transformers import AutoImageProcessor, SiglipForImageClassification | |
| from PIL import Image | |
| import io | |
| import os | |
| import tempfile | |
| from src.models.food_classification_model import FoodClassificationModel | |
| class PrithivMlFood101(FoodClassificationModel): | |
| """ | |
| Interface for accessing the PrithivML Food-101 model architecture. | |
| This model was already trained on the Food-101 dataset and performs well on it. | |
| It is supposed to serve as a benchmark for our own model finetuning and potentially | |
| as an alternative to be deployed. | |
| See it on Huggingface: https://huggingface.co/prithivMLmods/Food-101-93M. | |
| """ | |
| def __init__(self, model_name: str = "prithivMLmods/Food-101-93M"): | |
| """ | |
| Always load from the Hugging Face Hub. No local model storage. | |
| """ | |
| # Set up proper cache directory for HF Spaces (safe no-op locally) | |
| if not os.environ.get("HF_HOME"): | |
| cache_dir = tempfile.mkdtemp(prefix="transformers_cache_") | |
| os.environ["HF_HOME"] = str(cache_dir) | |
| cache_dir = os.environ.get("HF_HOME") | |
| # Load from the Hub | |
| self.model = SiglipForImageClassification.from_pretrained( | |
| model_name, | |
| cache_dir=cache_dir, | |
| ) | |
| self.processor = AutoImageProcessor.from_pretrained( | |
| model_name, | |
| cache_dir=cache_dir, | |
| use_fast=True, | |
| ) | |
| self.model_name = model_name | |
| self.model_path = model_name | |
| def classify(self, image: bytes) -> int: | |
| pil_image = Image.open(io.BytesIO(image)).convert("RGB") | |
| inputs = self.processor(images=pil_image, return_tensors="pt") | |
| with torch.no_grad(): | |
| outputs = self.model(**inputs) | |
| logits = outputs.logits | |
| probs = torch.nn.functional.softmax(logits, dim=1).squeeze() | |
| predicted_idx = torch.argmax(probs).item() | |
| return predicted_idx | |