Commit

app.py

@@ -1,63 +1,20 @@
-# Load Image to Text model
 import streamlit as st
-import torch
-import spaces
-from transformers import AutoProcessor, AutoModelForCausalLM, MBart50TokenizerFast, MBartForConditionalGeneration
-import requests
-# Carregamento de imagens locais
-import sys
-import cv2
+from transformers import pipeline
 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'
-
-tokenizer = MBart50TokenizerFast.from_pretrained(ckpt)
-translation_model = MBartForConditionalGeneration.from_pretrained(ckpt)
-
-tokenizer.src_lang = 'en_XX'
+st.title("Hot Dog? Or Not?")
 
 file_name = st.file_uploader("Upload a hot dog candidate image")
 
-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
-
-
-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'
-
-image = Image.open(file_name)
-# image = Image.open(requests.get(img_url, stream=True).raw)
+if file_name is not None:
+    col1, col2 = st.columns(2)
 
-caption = generate_caption(image)
-print(caption)
+    image = Image.open(file_name)
+    col1.image(image, use_column_width=True)
+    predictions = pipeline(image)
 
-translated_caption = translate(caption)
-print(translated_caption)
+    col2.header("Probabilities")
+    for p in predictions:
+        col2.subheader(f"{ p['label'] }: { round(p['score'] * 100, 1)}%")