mrrandom123 commited on
Commit
c5d6dd7
1 Parent(s): b3be49b
Files changed (1) hide show
  1. app.py +41 -0
app.py ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from PIL import Image
3
+ from transformers import BlipProcessor, BlipForConditionalGeneration, AutoTokenizer
4
+ import itertools
5
+ from nltk.corpus import stopwords
6
+ import nltk
7
+ import easyocr
8
+ import torch
9
+ import numpy as np
10
+ nltk.download('stopwords')
11
+
12
+ processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
13
+ model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
14
+
15
+ tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
16
+ reader = easyocr.Reader(['en'])
17
+ # set up Streamlit app
18
+ st.set_page_config(layout='wide', page_title='Image Hashtag Recommender')
19
+
20
+ def genrate_caption(image_file):
21
+ image = Image.open(image_file).convert('RGB')
22
+ inputs = processor(image, return_tensors="pt")
23
+ output_ids = model.generate(**inputs)
24
+ output_text = processor.decode(output_ids[0], skip_special_tokens=True)
25
+ return output_text
26
+
27
+ st.title("Image Caption and HashTag Recommender")
28
+ image_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
29
+
30
+ if image_file is not None:
31
+ try:
32
+ caption = genrate_caption(image_file)
33
+ if len(caption) > 0:
34
+ st.write(f"Caption : {caption}")
35
+
36
+ else:
37
+ st.write("No caption found for this image.")
38
+ except Exception as e:
39
+ st.write(f"Error: {e}")
40
+
41
+