import os
import gradio as gr
from PIL import Image
from lang_list import LANGS

##Image Classification
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
extractor = AutoFeatureExtractor.from_pretrained("rajistics/finetuned-indian-food")
model = AutoModelForImageClassification.from_pretrained("rajistics/finetuned-indian-food")

def image_to_text(imagepic):
  inputs = extractor(images=imagepic, return_tensors="pt")
  outputs = model(**inputs)
  logits = outputs.logits
  predicted_class_idx = logits.argmax(-1).item()
  return (model.config.id2label[predicted_class_idx])

##Translation
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
#Get list of language codes: https://github.com/facebookresearch/flores/tree/main/flores200#languages-in-flores-200
modelt = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
tokenizert = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")

def translation(text,target):
  translator = pipeline('translation', model=modelt, tokenizer=tokenizert, src_lang="eng_Latn", tgt_lang=target)
  output = translator(text)
  return (output[0]['translation_text'])

##Translation
demo = gr.Blocks()
with demo:
    image_file = gr.Image(type="pil")
    examples = gr.Examples(examples=[["003.jpg"],["126.jpg"],["401.jpg"]],inputs=[image_file])
    b1 = gr.Button("Recognize Image")
    text = gr.Textbox()
    b1.click(image_to_text, inputs=image_file, outputs=text)
    target = gr.Dropdown(LANGS,interactive=True,label="Target Language")
    b2 = gr.Button("Translation")
    out1 = gr.Textbox()
    b2.click(translation, inputs=[text,target], outputs=out1)
    #examples = gr.Examples(examples=[["003.jpg"]],inputs=[image_file])
demo.launch()