project / app.py
izhan001's picture
Create app.py
e06a9ce verified
import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification, AutoTokenizer, AutoModelForSeq2SeqLM
from datasets import load_dataset
from sklearn.model_selection import train_test_split
import torch
from PIL import Image
from torch.utils.data import Dataset
# Step 1: Load the World Cuisines dataset
ds = load_dataset("worldcuisines/food-kb")
# Access the 'main' dataset
dataset = ds['main']
# Check the structure of the dataset
print(dataset)
# Converting dataset to a list of dictionaries for easier manipulation
data_list = dataset.to_dict()['image1'] # Accessing the first image column (you can access others like image2, etc.)
# Now split the dataset into train and test
train_data, test_data = train_test_split(data_list, test_size=0.2)
# Check the shapes of train_data and test_data
print(f"Training data size: {len(train_data)}")
print(f"Testing data size: {len(test_data)}")
# Define a custom dataset class for the image classification task
class FoodDataset(Dataset):
def __init__(self, dataset, processor, max_length=256):
self.dataset = dataset
self.processor = processor
self.max_length = max_length
def __len__(self):
return len(self.dataset)
def __getitem__(self, idx):
item = self.dataset[idx]
# For simplicity, let's use image1 for training and test
image = Image.open(item['image1']) # Assuming 'image1' has the food images
label = item['fine_categories'] # You can modify this based on the label
# Process the image
encoding = self.processor(images=image, return_tensors="pt", padding=True, truncation=True)
# Return the input and target labels
return {
'input_ids': encoding['input_ids'].squeeze(),
'attention_mask': encoding['attention_mask'].squeeze(),
'labels': label # Assuming that 'fine_categories' is used as labels
}
# Step 2: Load the ViT model for image classification
processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
vit_model = AutoModelForImageClassification.from_pretrained("google/vit-base-patch16-224")
# Step 3: Load the text generation model (Gemini) for nutrition breakdown and diet plan
tokenizer = AutoTokenizer.from_pretrained("describeai/gemini")
gemini_model = AutoModelForSeq2SeqLM.from_pretrained("describeai/gemini")
# Helper function to get nutritional breakdown and allergen information
def get_nutrition_and_allergens(food_name):
# Look for the food item in the dataset
result = None
try:
dataset = ds['main'] # Access the correct dataset split
for item in dataset:
if food_name.lower() in item['name'].lower():
result = item
break
if result:
nutrition_info = result.get('nutrition', 'Nutrition information not available')
allergens = result.get('allergens', 'Allergen information not available')
diet_plan = f"This item is suitable for a diet including {result.get('suitable_for', 'N/A')}."
else:
nutrition_info = "Food item not found in the database."
allergens = "Allergen information not available."
diet_plan = "Diet plan not available for this food item."
except KeyError as e:
nutrition_info = f"Key error: {e}"
allergens = "Allergen information not available."
diet_plan = "Diet plan not available."
except Exception as e:
nutrition_info = f"An error occurred: {str(e)}"
allergens = "Allergen information not available."
diet_plan = "Diet plan not available."
return nutrition_info, allergens, diet_plan
# Main prediction function for the image classification and text generation
def predict(image):
try:
# Step 1: Classify the food item in the image using ViT model
inputs = processor(images=image, return_tensors="pt")
outputs = vit_model(**inputs)
# Get the predicted label (food item)
predicted_label = outputs.logits.argmax(-1).item()
# Get the food name from the class labels (assuming the model has the food labels)
class_labels = vit_model.config.id2label # Get the class label mapping
food_item = class_labels[predicted_label]
# Step 2: Generate nutritional breakdown, allergens, and diet plan
nutrition_info, allergens, diet_plan = get_nutrition_and_allergens(food_item)
# Step 3: Generate a detailed description using the Gemini model
description_input = f"Nutritional breakdown and diet plan for {food_item}"
diet_plan_text = tokenizer(description_input, return_tensors="pt", padding=True, truncation=True)
diet_plan_output = gemini_model.generate(**diet_plan_text)
diet_plan_text = tokenizer.decode(diet_plan_output[0], skip_special_tokens=True)
# Combine results into a single output
response = f"**Detected Food:** {food_item}\n\n"
response += f"**Nutrition Info:** {nutrition_info}\n\n"
response += f"**Allergens:** {allergens}\n\n"
response += f"**Diet Plan:** {diet_plan}\n\n"
response += f"**Detailed Diet Plan and Breakdown:** {diet_plan_text}"
except Exception as e:
response = f"Error: {str(e)}"
return response
# Step 4: Gradio Interface
interface = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs="text",
title="NutriScan: AI-Powered Food Analyzer",
description="Upload an image of food, and get a nutritional breakdown, allergen information, and diet plan recommendations.",
examples=[["path_to_example_image.jpg"]] # replace with paths to example images if needed
)
# Launch the Gradio interface
if __name__ == "__main__":
interface.launch()