app.py

#!/usr/bin/env python3
"""
Hugging Face Spaces FastAPI Food Recognition Service
Optimized for Hugging Face Spaces deployment
"""

import gradio as gr
import requests
import base64
import io
from PIL import Image
import torch
from transformers import pipeline
import logging
from datetime import datetime
import os

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global variables for model
classifier = None
model_loaded = False

# Model configuration
MODEL_ID = "BinhQuocNguyen/food-recognition-vit"
FOOD_CLASSES = [
    "apple_pie", "caesar_salad", "chocolate_cake", "cup_cakes", "donuts",
    "hamburger", "ice_cream", "pancakes", "pizza", "waffles"
]

def load_model():
    """Load the Hugging Face model"""
    global classifier, model_loaded
    try:
        logger.info(f"Loading model: {MODEL_ID}")
        classifier = pipeline(
            "image-classification",
            model=MODEL_ID,
            device=-1,  # Use CPU (change to 0 for GPU)
            use_fast=True  # Use fast image processor
        )
        model_loaded = True
        logger.info("Model loaded successfully!")
        return True
    except Exception as e:
        logger.error(f"Failed to load model: {e}")
        model_loaded = False
        return False

def preprocess_image(image):
    """Preprocess uploaded image"""
    try:
        if isinstance(image, str):
            # If it's a file path
            image = Image.open(image)
        elif hasattr(image, 'convert'):
            # If it's already a PIL Image
            pass
        else:
            # If it's numpy array or other format
            image = Image.fromarray(image)
        
        # Convert to RGB if necessary
        if image.mode != 'RGB':
            image = image.convert('RGB')
        return image
    except Exception as e:
        raise ValueError(f"Invalid image format: {e}")

def predict_food(image):
    """Predict food type from image"""
    if not model_loaded:
        return "Model not loaded. Please try again.", None
    
    try:
        # Preprocess image
        processed_image = preprocess_image(image)
        
        # Make prediction
        results = classifier(processed_image)
        
        # Format results
        predictions = []
        for result in results:
            predictions.append({
                'label': result['label'],
                'confidence': result['score']
            })
        
        # Get top prediction
        top_prediction = predictions[0]
        confidence_percent = top_prediction['confidence'] * 100
        
        # Create result text
        result_text = f"🍕 **Predicted Food:** {top_prediction['label'].replace('_', ' ').title()}\n"
        result_text += f"🎯 **Confidence:** {confidence_percent:.1f}%\n\n"
        result_text += "**Top 3 Predictions:**\n"
        
        for i, pred in enumerate(predictions[:3], 1):
            food_name = pred['label'].replace('_', ' ').title()
            conf_percent = pred['confidence'] * 100
            result_text += f"{i}. {food_name}: {conf_percent:.1f}%\n"
        
        return result_text, processed_image
        
    except Exception as e:
        logger.error(f"Prediction error: {e}")
        return f"❌ Error: {str(e)}", None

def get_model_info():
    """Get model information"""
    return {
        "model_id": MODEL_ID,
        "model_url": f"https://huggingface.co/{MODEL_ID}",
        "classes": FOOD_CLASSES,
        "num_classes": len(FOOD_CLASSES),
        "device": "cpu"
    }

# Load model on startup
load_model()

# Create Gradio interface
def create_interface():
    """Create the Gradio interface"""
    
    with gr.Blocks(
        title="Food Recognition API",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 800px !important;
            margin: auto !important;
        }
        """
    ) as interface:
        
        gr.Markdown("""
        # 🍕 Food Recognition API
        
        Upload an image of food and get instant predictions! This API uses a Vision Transformer model 
        trained to recognize 10 different types of food.
        
        **Supported Food Types:** Apple Pie, Caesar Salad, Chocolate Cake, Cup Cakes, Donuts, 
        Hamburger, Ice Cream, Pancakes, Pizza, Waffles
        
        **How to use:** Simply drag and drop an image or click to upload, then click "Predict Food"!
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(
                    label="Upload Food Image",
                    type="pil",
                    height=300
                )
                
                predict_btn = gr.Button(
                    "🔍 Predict Food",
                    variant="primary",
                    size="lg"
                )
                
                gr.Markdown("""
                ### 📊 Model Information
                - **Model:** Vision Transformer (ViT)
                - **Accuracy:** 68%
                - **Classes:** 10 food types
                - **Source:** [Hugging Face Model](https://huggingface.co/BinhQuocNguyen/food-recognition-vit)
                """)
            
            with gr.Column(scale=1):
                output_text = gr.Markdown(
                    label="Prediction Results",
                    value="👆 Upload an image and click 'Predict Food' to get started!"
                )
                
                output_image = gr.Image(
                    label="Processed Image",
                    height=300
                )
        
        
        # Event handlers
        predict_btn.click(
            fn=predict_food,
            inputs=image_input,
            outputs=[output_text, output_image]
        )
        
        # Footer
        gr.Markdown("""
        ---
        **Built with:** FastAPI, Gradio, Hugging Face Transformers, PyTorch
        
        **Model Performance:** 68% accuracy on 10 food classes
        
        **API Endpoints:** Available at `/docs` for programmatic access
        """)
    
    return interface

# Create the interface
interface = create_interface()

# FastAPI app for additional endpoints
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional
import uvicorn

# Initialize FastAPI app
app = FastAPI(
    title="Food Recognition API",
    description="API for food recognition using Hugging Face Vision Transformer model",
    version="1.0.0"
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Pydantic models
class PredictionResult(BaseModel):
    label: str
    confidence: float

class PredictionResponse(BaseModel):
    predictions: List[PredictionResult]
    processing_time: float
    model_info: dict

class HealthResponse(BaseModel):
    status: str
    model_loaded: bool
    timestamp: str
    model_info: Optional[dict] = None

# FastAPI routes
@app.get("/api")
async def api_info():
    """API information endpoint"""
    return {
        "message": "Food Recognition API",
        "version": "1.0.0",
        "model": MODEL_ID,
        "gradio_interface": "/",
        "api_docs": "/docs"
    }

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Health check endpoint"""
    return HealthResponse(
        status="healthy" if model_loaded else "unhealthy",
        model_loaded=model_loaded,
        timestamp=datetime.now().isoformat(),
        model_info=get_model_info() if model_loaded else None
    )

@app.get("/classes")
async def get_classes():
    """Get supported food classes"""
    return {
        "classes": FOOD_CLASSES,
        "num_classes": len(FOOD_CLASSES),
        "model_id": MODEL_ID
    }

@app.get("/model-info")
async def get_model_information():
    """Get detailed model information"""
    if not model_loaded:
        return {"error": "Model not loaded"}
    return get_model_info()

# Mount Gradio interface
app = gr.mount_gradio_app(app, interface, path="/")

if __name__ == "__main__":
    # For local development
    uvicorn.run(app, host="0.0.0.0", port=7860)