Commit

app.py

@@ -1,35 +1,63 @@
+# Load Image to Text model
 import streamlit as st
 import torch
 import spaces
-from transformers import pipeline, AutoModelForCausalLM, AutoProcessor
+from transformers import AutoProcessor, AutoModelForCausalLM, MBart50TokenizerFast, MBartForConditionalGeneration
+import requests
+# Carregamento de imagens locais
+import sys
+import cv2
 from PIL import Image
+# Load Translation model
 
+image_processor = AutoProcessor.from_pretrained("sezenkarakus/image-GIT-description-model-v3")
+image_to_text_model = AutoModelForCausalLM.from_pretrained("sezenkarakus/image-GIT-description-model-v3")
 
-pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")
+ckpt = 'Narrativa/mbart-large-50-finetuned-opus-en-pt-translation'
 
-st.title("Hot Dog? Or Not?")
+tokenizer = MBart50TokenizerFast.from_pretrained(ckpt)
+translation_model = MBartForConditionalGeneration.from_pretrained(ckpt)
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-image_model_id = "microsoft/git-large-coco"
-image_processor = AutoProcessor.from_pretrained(image_model_id)
-image_model = AutoModelForCausalLM.from_pretrained(image_model_id).to(device)
+tokenizer.src_lang = 'en_XX'
 
 file_name = st.file_uploader("Upload a hot dog candidate image")
 
-if file_name is not None:
-    col1, col2 = st.columns(2)
+def generate_caption(image):
+  pixel_values = image_processor(images=image, return_tensors="pt").pixel_values
+  generated_ids = image_to_text_model.generate(pixel_values=pixel_values, max_length=200)
+  generated_caption = image_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+  return generated_caption
+
+def translate(text):
+  inputs = tokenizer(text, return_tensors='pt')
+  input_ids = inputs.input_ids
+  attention_mask = inputs.attention_mask
+
+  try:
+    input_ids = input_ids.to('cuda')
+    attention_mask = attention_mask.to('cuda')
+    model = translation_model.to("cuda")
+  except:
+    print('No NVidia GPU, model performance may not be as good')
+    model = translation_model
+
+  output = model.generate(input_ids, attention_mask=attention_mask, forced_bos_token_id=tokenizer.lang_code_to_id['pt_XX'])
+  translated = tokenizer.decode(output[0], skip_special_tokens=True)
+
+  return translated
+
 
-    image = Image.open(file_name)
-    col1.image(image, use_column_width=True)
-    predictions = pipeline(image)
+img_url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
+# img_url = 'https://farm4.staticflickr.com/3733/9000662079_ce3599d0d8_z.jpg'
+# img_url = 'https://farm4.staticflickr.com/3088/5793281956_2a15b2559c_z.jpg'
+# img_url = 'https://farm5.staticflickr.com/4073/4816939054_844feb0078_z.jpg'
 
-    col2.header("Probabilities")
-    for p in predictions:
-        col2.subheader(f"{ p['label'] }: { round(p['score'] * 100, 1)}%")
+image = Image.open(file_name)
+# image = Image.open(requests.get(img_url, stream=True).raw)
 
-pixel_values = image_processor(images=image, return_tensors="pt").pixel_values
+caption = generate_caption(image)
+print(caption)
 
-generated_ids = image_model.generate(pixel_values=pixel_values, max_length=50)
-generated_caption = image_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-print(generated_caption)
+translated_caption = translate(caption)
+print(translated_caption)