jframed281/ImageToTextGenerator

1# -*- coding: utf-8 -*-2"""imagetortext.ipynb3 4Automatically generated by Colaboratory.5 6Original file is located at7    https://colab.research.google.com/drive/1UXh8tivt-4vHaBeXgfyq-TYLvpJAMZVV8"""9 10!pip install transformers dataset11!pip install -q gradio12!pip install sentencepiece13!pip install googletrans==3.1.0a014 15import gradio as gr16import requests17from PIL import Image18from torchvision import transforms19from transformers import pipeline20import torch21import sentencepiece22import re23import googletrans 24from googletrans import Translator25 26model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True).eval()27 28from transformers import SegformerFeatureExtractor, SegformerForImageClassification,  T5Tokenizer, T5Model29from PIL import Image30import requests31 32 33def loadImageToText(image, argument):34  # url = "https://media.istockphoto.com/id/470604022/es/foto/%C3%A1rbol-de-manzano.jpg?s=1024x1024&w=is&k=20&c=R7b6jPeTGsDw75Sqn3VwpNRckqlAkJNPLelb48pCk2U="35  # image = Image.open(requests.get(url, stream=True).raw)36  feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/mit-b2")37  model = SegformerForImageClassification.from_pretrained("nvidia/mit-b2")38  translator = Translator()39 40  inputs = feature_extractor(images=image, return_tensors="pt")41  outputs = model(**inputs)42  logits = outputs.logits43  # model predicts one of the 1000 ImageNet classes44  predicted_class_idx = logits.argmax(-1).item()45  part_args = f"<"+re.sub("[^(\w|<|>)]+(?=\w)", "><", argument) + ">"46  story_gen = pipeline("text-generation", "pranavpsv/gpt2-genre-story-generator")47  story_text = story_gen(part_args + model.config.id2label[predicted_class_idx])48 49  generate_text_stroy = story_text[0]["generated_text"]50  ln_text_story = generate_text_stroy[len(part_args):len(generate_text_stroy)]51 52  translated_ita = translator.translate(ln_text_story, src='en', dest='es')53 54  return translated_ita.text55 56 57 58#print(loadImageToText("animal,super")) # borra el argumento 'image' y sus variables internas si quieres probarlo desde aquí.59 60gr.Interface(fn=loadImageToText,61             inputs=[gr.Image(), gr.Text(label="Argumentos base", placeholder="Verano, película,playa, superhéroe, animal")],62             outputs="text").launch()