Spaces:

calcuplate
/

ingredientClassificationModel

Sleeping

Update app.py

aec5a02 verified 12 months ago

5.18 kB

import gradio as gr
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import json

# Define class names
class_names = ['acorn_squash', 'almond', 'almonds', 'anchovy_(fish)', 'apple', 'apricot', 'artichoke', 'arugula', 'asparagus', 'avocado', 'baguette', 'banana', 'barley', 'barley_(grain)', 'beef_(meat)', 'beet', 'black_beans', 'black_pepper_(spice)', 'blackberry', 'bok_choy', 'bread_(loaf)', 'breadcrumbs', 'broccoli', 'brussels_sprouts', 'butter_(dairy)', 'butternut_squash', 'cabbage', 'canola_oil', 'cantaloupe', 'carrot', 'cashews', 'cauliflower', 'celery', 'cheddar_cheese', 'cherry', 'chicken_(meat)', 'chickpeas', 'chive', 'chocolate_chips', 'clams', 'clams_(seafood)', 'cocoa_powder', 'coconut', 'cod_(fish)', 'condensed_milk', 'confectioners_sugar', 'corn', 'corn_syrup', 'cornflakes', 'cornmeal', 'cottage_cheese', 'crab_(seafood)', 'crackers', 'cranberry', 'cream_(dairy)', 'cream_cheese', 'cucumber', 'date_(fruit)', 'dragonfruit', 'duck_(meat)', 'egg', 'eggplant', 'evaporated_milk', 'feta', 'feta_cheese', 'fig', 'fish_sauce', 'garlic', 'goat_cheese', 'grape', 'ground_beef', 'ground_pork', 'ground_turkey', 'guava', 'honeydew', 'jackfruit', 'kale', 'ketchup', 'kidney_beans', 'kiwi_(fruit)', 'lamb_(meat)', 'leek', 'lemon_(fruit)', 'lentils', 'lettuce', 'lime_(fruit)', 'lobster_(seafood)', 'lychee', 'mango', 'mayonnaise', 'meatballs', 'milk_(dairy)', 'molasses', 'mozzarella_cheese', 'mulberry', 'mushroom', 'mussels_(seafood)', 'mustard_(condiment)', 'mustard_greens', 'navy_beans', 'nectarine', 'noodles_(cooked)', 'oats', 'oats_(grain)', 'octopus_(seafood)', 'okra', 'olive', 'olive_oil', 'onion', 'orange_(fruit)', 'oyster_sauce', 'papaya', 'parmesan_cheese', 'parsnip', 'passionfruit', 'pasta_(cooked)', 'peach', 'peanut', 'peanut_butter', 'pear', 'pecans', 'pepper', 'persimmon', 'pineapple', 'pinto_beans', 'pita_bread', 'plum', 'pomegranate', 'pork_(meat)', 'potato', 'powdered_milk', 'powdered_sugar', 'pumpkin_seeds', 'quinoa', 'quinoa_(grain)', 'radish', 'raspberry', 'rice_(brown,_grain)', 'rice_(white,_grain)', 'ricotta', 'ricotta_cheese', 'rolled_oats', 'salmon', 'salmon_(fish)', 'salt', 'sardine_(fish)', 'scallion', 'scallops_(seafood)', 'seitan', 'sesame_oil', 'sesame_seeds', 'shallot', 'shrimp_(seafood)', 'sour_cream', 'soy_sauce', 'spinach', 'split_peas', 'squid_(seafood)', 'starfruit', 'strawberry', 'sunflower_oil', 'sunflower_seeds', 'sweet_potato', 'swiss_chard', 'tangerine', 'tempeh', 'tofu', 'tomato', 'tortilla_(flatbread)', 'tortillas', 'tuna', 'tuna_(fish)', 'turkey_(meat)', 'turnip', 'vegetable_oil', 'walnut', 'walnuts', 'watermelon', 'wheat_flour', 'whipping_cream', 'yam', 'yogurt_(dairy)', 'zucchini']

# Load model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

model = models.mobilenet_v2(pretrained=False)
num_ftrs = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_ftrs, len(class_names))
model.load_state_dict(torch.load("ingredient_recognition_model.pth", map_location=device))
model = model.to(device)
model.eval()

# Image transformation
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

def clean_ingredient_name(name):
    """Clean up ingredient name for display"""
    name = name.split('_(')[0]
    name = name.replace('_', ' ')
    return name.title()

def predict(image):
    """Predict ingredients from image"""
    if image is None:
        return {"error": "No image provided"}
    
    try:
        # Convert to PIL Image if needed
        if not isinstance(image, Image.Image):
            image = Image.fromarray(image)
        
        # Convert to RGB
        image = image.convert("RGB")
        
        # Transform image
        input_tensor = transform(image).unsqueeze(0).to(device)
        
        # Run inference
        with torch.no_grad():
            outputs = model(input_tensor)
            probs = torch.nn.functional.softmax(outputs[0], dim=0)
            
            # Get top 5 predictions
            k = min(5, len(class_names))
            top_probs, top_idxs = torch.topk(probs, k)
        
        # Build results dictionary
        results = {}
        for prob, idx in zip(top_probs, top_idxs):
            raw_name = class_names[idx]
            clean_name = clean_ingredient_name(raw_name)
            confidence = prob.item()
            results[clean_name] = float(confidence)
        
        return results
        
    except Exception as e:
        return {"error": str(e)}

# Create Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Food Image"),
    outputs=gr.Label(num_top_classes=5, label="Predictions"),
    title="Ingredient Classification Model",
    description="Upload an image of food to identify the ingredients. This model recognizes over 200 different foods and ingredients.",
    examples=[
        # Add example images if you have them
    ],
    theme=gr.themes.Soft(),
    allow_flagging="never"
)

if __name__ == "__main__":
    iface.launch()

	import gradio as gr
	import torch
	import torch.nn as nn
	from torchvision import models, transforms
	from PIL import Image
	import json

	# Define class names
	class_names = ['acorn_squash', 'almond', 'almonds', 'anchovy_(fish)', 'apple', 'apricot', 'artichoke', 'arugula', 'asparagus', 'avocado', 'baguette', 'banana', 'barley', 'barley_(grain)', 'beef_(meat)', 'beet', 'black_beans', 'black_pepper_(spice)', 'blackberry', 'bok_choy', 'bread_(loaf)', 'breadcrumbs', 'broccoli', 'brussels_sprouts', 'butter_(dairy)', 'butternut_squash', 'cabbage', 'canola_oil', 'cantaloupe', 'carrot', 'cashews', 'cauliflower', 'celery', 'cheddar_cheese', 'cherry', 'chicken_(meat)', 'chickpeas', 'chive', 'chocolate_chips', 'clams', 'clams_(seafood)', 'cocoa_powder', 'coconut', 'cod_(fish)', 'condensed_milk', 'confectioners_sugar', 'corn', 'corn_syrup', 'cornflakes', 'cornmeal', 'cottage_cheese', 'crab_(seafood)', 'crackers', 'cranberry', 'cream_(dairy)', 'cream_cheese', 'cucumber', 'date_(fruit)', 'dragonfruit', 'duck_(meat)', 'egg', 'eggplant', 'evaporated_milk', 'feta', 'feta_cheese', 'fig', 'fish_sauce', 'garlic', 'goat_cheese', 'grape', 'ground_beef', 'ground_pork', 'ground_turkey', 'guava', 'honeydew', 'jackfruit', 'kale', 'ketchup', 'kidney_beans', 'kiwi_(fruit)', 'lamb_(meat)', 'leek', 'lemon_(fruit)', 'lentils', 'lettuce', 'lime_(fruit)', 'lobster_(seafood)', 'lychee', 'mango', 'mayonnaise', 'meatballs', 'milk_(dairy)', 'molasses', 'mozzarella_cheese', 'mulberry', 'mushroom', 'mussels_(seafood)', 'mustard_(condiment)', 'mustard_greens', 'navy_beans', 'nectarine', 'noodles_(cooked)', 'oats', 'oats_(grain)', 'octopus_(seafood)', 'okra', 'olive', 'olive_oil', 'onion', 'orange_(fruit)', 'oyster_sauce', 'papaya', 'parmesan_cheese', 'parsnip', 'passionfruit', 'pasta_(cooked)', 'peach', 'peanut', 'peanut_butter', 'pear', 'pecans', 'pepper', 'persimmon', 'pineapple', 'pinto_beans', 'pita_bread', 'plum', 'pomegranate', 'pork_(meat)', 'potato', 'powdered_milk', 'powdered_sugar', 'pumpkin_seeds', 'quinoa', 'quinoa_(grain)', 'radish', 'raspberry', 'rice_(brown,_grain)', 'rice_(white,_grain)', 'ricotta', 'ricotta_cheese', 'rolled_oats', 'salmon', 'salmon_(fish)', 'salt', 'sardine_(fish)', 'scallion', 'scallops_(seafood)', 'seitan', 'sesame_oil', 'sesame_seeds', 'shallot', 'shrimp_(seafood)', 'sour_cream', 'soy_sauce', 'spinach', 'split_peas', 'squid_(seafood)', 'starfruit', 'strawberry', 'sunflower_oil', 'sunflower_seeds', 'sweet_potato', 'swiss_chard', 'tangerine', 'tempeh', 'tofu', 'tomato', 'tortilla_(flatbread)', 'tortillas', 'tuna', 'tuna_(fish)', 'turkey_(meat)', 'turnip', 'vegetable_oil', 'walnut', 'walnuts', 'watermelon', 'wheat_flour', 'whipping_cream', 'yam', 'yogurt_(dairy)', 'zucchini']

	# Load model
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	print(f"Using device: {device}")

	model = models.mobilenet_v2(pretrained=False)
	num_ftrs = model.classifier[1].in_features
	model.classifier[1] = nn.Linear(num_ftrs, len(class_names))
	model.load_state_dict(torch.load("ingredient_recognition_model.pth", map_location=device))
	model = model.to(device)
	model.eval()

	# Image transformation
	transform = transforms.Compose([
	transforms.Resize((224, 224)),
	transforms.ToTensor(),
	transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
	])

	def clean_ingredient_name(name):
	"""Clean up ingredient name for display"""
	name = name.split('_(')[0]
	name = name.replace('_', ' ')
	return name.title()

	def predict(image):
	"""Predict ingredients from image"""
	if image is None:
	return {"error": "No image provided"}

	try:
	# Convert to PIL Image if needed
	if not isinstance(image, Image.Image):
	image = Image.fromarray(image)

	# Convert to RGB
	image = image.convert("RGB")

	# Transform image
	input_tensor = transform(image).unsqueeze(0).to(device)

	# Run inference
	with torch.no_grad():
	outputs = model(input_tensor)
	probs = torch.nn.functional.softmax(outputs[0], dim=0)

	# Get top 5 predictions
	k = min(5, len(class_names))
	top_probs, top_idxs = torch.topk(probs, k)

	# Build results dictionary
	results = {}
	for prob, idx in zip(top_probs, top_idxs):
	raw_name = class_names[idx]
	clean_name = clean_ingredient_name(raw_name)
	confidence = prob.item()
	results[clean_name] = float(confidence)

	return results

	except Exception as e:
	return {"error": str(e)}

	# Create Gradio interface
	iface = gr.Interface(
	fn=predict,
	inputs=gr.Image(type="pil", label="Upload Food Image"),
	outputs=gr.Label(num_top_classes=5, label="Predictions"),
	title="Ingredient Classification Model",
	description="Upload an image of food to identify the ingredients. This model recognizes over 200 different foods and ingredients.",
	examples=[
	# Add example images if you have them
	],
	theme=gr.themes.Soft(),
	allow_flagging="never"
	)

	if __name__ == "__main__":
	iface.launch()