File size: 2,006 Bytes
3d429a8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
from transformers import AutoProcessor, AutoModelForImageClassification, pipeline
from PIL import Image
import torch
def load_classification_model():
processor = AutoProcessor.from_pretrained("Shresthadev403/food-image-classification")
model = AutoModelForImageClassification.from_pretrained("Shresthadev403/food-image-classification")
return processor, model
def load_text_generator():
return pipeline( "text2text-generation",model="distilgpt2")
processor, model = load_classification_model()
text_generator = load_text_generator()
def predict_dish(image: Image.Image):
try:
image = Image.open(image).convert("RGB")
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
predicted_class_idx = logits.argmax(-1).item()
label = model.config.id2label[predicted_class_idx]
return label.lower().replace(" ", "_")
except Exception as e:
print(f"❌ Error in dish prediction: {e}")
return "unknown_dish"
def generate_recipe(dish, diet=None, cuisine=None, cook_time=None):
filters = []
if diet and diet != "Any":
filters.append(f"{diet} diet")
if cuisine and cuisine != "Any":
filters.append(f"{cuisine} cuisine")
if cook_time and cook_time != "Any":
filters.append(f"ready in {cook_time}")
filter_text = ", ".join(filters)
prompt = f"""
Create a step-by-step recipe for {dish}.
Include:
- Ingredients with quantities
- Step-by-step instructions cooking steps
Make sure it's a {filter_text} recipe."""
try:
result = text_generator(prompt.strip(), max_length=282, do_sample=False)
return result[0]['generated_text']
except Exception as e:
print(f"❌ Error generating recipe: {e}")
return "Sorry, couldn't generate a recipe at the moment."
|