from transformers import ViTConfig, ViTForImageClassification
from transformers import ViTFeatureExtractor
from PIL import Image
import requests
import matplotlib.pyplot as plt
import gradio as gr
from gradio.mix import Parallel
from transformers import ImageClassificationPipeline, PerceiverForImageClassificationConvProcessing, PerceiverFeatureExtractor
from transformers import VisionEncoderDecoderModel
from transformers import AutoTokenizer
import torch

# https://github.com/NielsRogge/Transformers-Tutorials/blob/master/HuggingFace_vision_ecosystem_overview_(June_2022).ipynb
# option 1: load with randomly initialized weights (train from scratch)

config = ViTConfig(num_hidden_layers=12, hidden_size=768)
model = ViTForImageClassification(config)

#print(config)

feature_extractor = ViTFeatureExtractor()
# or, to load one that corresponds to a checkpoint on the hub:
#feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")

#the following gets called by classify_image() 
feature_extractor = PerceiverFeatureExtractor.from_pretrained("deepmind/vision-perceiver-conv")
model = PerceiverForImageClassificationConvProcessing.from_pretrained("deepmind/vision-perceiver-conv")
#google/vit-base-patch16-224, deepmind/vision-perceiver-conv
image_pipe = ImageClassificationPipeline(model=model, feature_extractor=feature_extractor)


def self_caption(image):
  repo_name = "ydshieh/vit-gpt2-coco-en"
  #test_image = "cats.jpg"
  test_image = image
  #url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
  #test_image = Image.open(requests.get(url, stream=True).raw)
  #test_image.save("cats.png")
  
  feature_extractor2 = ViTFeatureExtractor.from_pretrained(repo_name)
  tokenizer = AutoTokenizer.from_pretrained(repo_name)
  model2 = VisionEncoderDecoderModel.from_pretrained(repo_name)
  pixel_values = feature_extractor2(test_image, return_tensors="pt").pixel_values
  print("Pixel Values")
  print(pixel_values)
  # autoregressively generate text (using beam search or other decoding strategy)
  generated_ids = model2.generate(pixel_values, max_length=16, num_beams=4, return_dict_in_generate=True)
  
  # decode into text
  preds = tokenizer.batch_decode(generated_ids[0], skip_special_tokens=True)
  preds = [pred.strip() for pred in preds]
  print("Predictions")
  print(preds)
  print("The preds type is : ",type(preds))
  pred_keys = ["Prediction"]
  pred_value = preds

  pred_dictionary = dict(zip(pred_keys, pred_value))
  print("Pred dictionary")
  print(pred_dictionary)
  #return(pred_dictionary)
  preds = ' '.join(preds)
  return preds


def classify_image(image):
  results = image_pipe(image)
  
  print("RESULTS")
  print(results)
  # convert to format Gradio expects
  output = {}
  for prediction in results:
    predicted_label = prediction['label']
    score = prediction['score']
    output[predicted_label] = score
  print("OUTPUT")
  print(output)
  return output


image = gr.inputs.Image(type="pil")
label = gr.outputs.Label(num_top_classes=5)
examples = [["cats.jpg"]]
title = "Generate a Story from an Image"
description = "Demo for classifying images with Perceiver IO. To use it, simply upload an image and click 'submit', a caption is autogenerated as well"
article = "<p style='text-align: center'></p>"

img_info1 = gr.Interface(
    fn=classify_image,
    inputs=image,
    outputs=label,
)

img_info2 = gr.Interface(
    fn=self_caption,
    inputs=image,
    #outputs=label,
    outputs = [
    gr.outputs.Textbox(label = 'Caption')
],
)

Parallel(img_info1,img_info2, inputs=image, title=title, description=description, examples=examples, enable_queue=True).launch(debug=True)
#Parallel(img_info1,img_info2, inputs=image, outputs=label, title=title, description=description, examples=examples, enable_queue=True).launch(debug=True)