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()