File size: 1,405 Bytes
e1396a6
92ff04c
8f8159e
2e2d57c
e1396a6
59b6d58
 
8f8159e
72c8a31
92ff04c
 
 
 
 
 
 
 
 
 
 
 
e1396a6
8f8159e
e1396a6
c023386
e894482
 
c023386
8f8159e
 
 
 
92ff04c
 
 
72c8a31
 
92ff04c
8f8159e
 
 
 
 
 
 
d253c2c
72c8a31
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
import streamlit as st
from transformers import pipeline, AutoModelWithLMHead, AutoTokenizer
from PIL import Image
import torch

st.set_page_config(layout="wide")

image_pipe = pipeline("image-classification")
text_pipe = pipeline("text-generation")
k2t_tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-common_gen")
k2t_model = AutoModelWithLMHead.from_pretrained("mrm8488/t5-base-finetuned-common_gen")

def gen_sentence(words, max_length=32):
  input_text = words
  features = k2t_tokenizer([input_text], return_tensors='pt')

  output = k2t_model.generate(input_ids=features['input_ids'], 
               attention_mask=features['attention_mask'],
               max_length=max_length)

  return k2t_tokenizer.decode(output[0], skip_special_tokens=True)

img = st.file_uploader(label='Upload jpg or png to create post',type=['jpg','png'])

if img is None:
	torch.hub.download_url_to_file('https://assets.epicurious.com/photos/5761d0268accf290434553aa/master/pass/panna-cotta.jpg', "img.jpg")
	img = "img.jpg"

with Image.open(img) as img:

	results = image_pipe(img)
	
	keywords = ""
	for keyword in results:
		keywords += keyword["label"].split(',')[0]
		
	post_text = text_pipe(gen_sentence(keywords))[0]["generated_text"]
	
	col1, col2 = st.columns(2)
	
	with col1:
		st.subheader("Your image")
		st.image(img)
	
	with col2:
		st.subheader("Generated text")
		st.write(post_text)