cookai / app.py
ajdarshaydullin's picture
Add application file
c725fc9
raw
history blame
2.04 kB
import gradio as gr
# Use a pipeline as a high-level helper
from transformers import pipeline
detector = pipeline("object-detection", model="hustvl/yolos-tiny")
# Load model directly
from transformers import AutoImageProcessor, AutoModelForImageClassification
processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50")
model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-50")
def food_classifier(image):
inputs = processor(image, return_tensors="pt")
logits = model(**inputs).logits
predicted_label = logits.argmax(-1).item()
label = model.config.id2label[predicted_label]
return [{'label': label, 'score': 0.0}]
# Use a pipeline as a high-level helper
from transformers import pipeline
food_classifier = pipeline("image-classification", model="facebook/deit-base-distilled-patch16-384")
def get_ingridients_list(image, score_threshold=.85):
objects = detector(image)
ingridients = []
for obj in objects:
cropped_image = image.crop((obj['box']['xmin'], obj['box']['ymin'], obj['box']['xmax'], obj['box']['ymax']))
classes = food_classifier(cropped_image)
best_match = max(classes, key=lambda x: x['score'])
if best_match['score'] > score_threshold:
ingridients.append(best_match['label'])
return list(set(ingridients))
def get_ingridients(image):
ingridients = get_ingridients_list(image)
return ', '.join(ingridients)
#text_to_text = pipeline("text-generation", model="ai-forever/mGPT")
def get_reciepe(ingridients):
return 'dish of ' + ingridients
def get_answer(image):
ingridients = get_ingridients(image)
return get_reciepe(ingridients)
# Create a Gradio interface
iface = gr.Interface(
fn=get_answer, # Function to call
inputs=gr.Image(label="Upload an image", type="pil"), # Input type: Image
outputs=gr.Markdown(label="Classification Result"), # Output type: Markdown
title="Food Ingredient Classifier",
description="Upload an image of a food ingredient to classify it."
)
# Launch the Gradio app
iface.launch()