seu-ebers commited on
Commit
595eae9
·
1 Parent(s): f28773b
Files changed (1) hide show
  1. app.py +11 -54
app.py CHANGED
@@ -1,63 +1,20 @@
1
- # Load Image to Text model
2
  import streamlit as st
3
- import torch
4
- import spaces
5
- from transformers import AutoProcessor, AutoModelForCausalLM, MBart50TokenizerFast, MBartForConditionalGeneration
6
- import requests
7
- # Carregamento de imagens locais
8
- import sys
9
- import cv2
10
  from PIL import Image
11
- # Load Translation model
12
 
13
- image_processor = AutoProcessor.from_pretrained("sezenkarakus/image-GIT-description-model-v3")
14
- image_to_text_model = AutoModelForCausalLM.from_pretrained("sezenkarakus/image-GIT-description-model-v3")
15
 
16
- ckpt = 'Narrativa/mbart-large-50-finetuned-opus-en-pt-translation'
17
-
18
- tokenizer = MBart50TokenizerFast.from_pretrained(ckpt)
19
- translation_model = MBartForConditionalGeneration.from_pretrained(ckpt)
20
-
21
- tokenizer.src_lang = 'en_XX'
22
 
23
  file_name = st.file_uploader("Upload a hot dog candidate image")
24
 
25
- def generate_caption(image):
26
- pixel_values = image_processor(images=image, return_tensors="pt").pixel_values
27
- generated_ids = image_to_text_model.generate(pixel_values=pixel_values, max_length=200)
28
- generated_caption = image_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
29
-
30
- return generated_caption
31
-
32
- def translate(text):
33
- inputs = tokenizer(text, return_tensors='pt')
34
- input_ids = inputs.input_ids
35
- attention_mask = inputs.attention_mask
36
-
37
- try:
38
- input_ids = input_ids.to('cuda')
39
- attention_mask = attention_mask.to('cuda')
40
- model = translation_model.to("cuda")
41
- except:
42
- print('No NVidia GPU, model performance may not be as good')
43
- model = translation_model
44
-
45
- output = model.generate(input_ids, attention_mask=attention_mask, forced_bos_token_id=tokenizer.lang_code_to_id['pt_XX'])
46
- translated = tokenizer.decode(output[0], skip_special_tokens=True)
47
-
48
- return translated
49
-
50
-
51
- img_url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
52
- # img_url = 'https://farm4.staticflickr.com/3733/9000662079_ce3599d0d8_z.jpg'
53
- # img_url = 'https://farm4.staticflickr.com/3088/5793281956_2a15b2559c_z.jpg'
54
- # img_url = 'https://farm5.staticflickr.com/4073/4816939054_844feb0078_z.jpg'
55
-
56
- image = Image.open(file_name)
57
- # image = Image.open(requests.get(img_url, stream=True).raw)
58
 
59
- caption = generate_caption(image)
60
- print(caption)
 
61
 
62
- translated_caption = translate(caption)
63
- print(translated_caption)
 
 
 
1
  import streamlit as st
2
+ from transformers import pipeline
 
 
 
 
 
 
3
  from PIL import Image
 
4
 
5
+ pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")
 
6
 
7
+ st.title("Hot Dog? Or Not?")
 
 
 
 
 
8
 
9
  file_name = st.file_uploader("Upload a hot dog candidate image")
10
 
11
+ if file_name is not None:
12
+ col1, col2 = st.columns(2)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
13
 
14
+ image = Image.open(file_name)
15
+ col1.image(image, use_column_width=True)
16
+ predictions = pipeline(image)
17
 
18
+ col2.header("Probabilities")
19
+ for p in predictions:
20
+ col2.subheader(f"{ p['label'] }: { round(p['score'] * 100, 1)}%")