Update pages/ImageToText.py

pages/ImageToText.py

@@ -11,33 +11,39 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 vitgpt_model.to(device)
 
-def generate_caption(processor, model, image, tokenizer=None):
+def generate_caption(processor, model, image, tokenizer=None, num_seq):
     
     inputs = processor(images=image, return_tensors="pt").to(device)
     generated_ids = model.generate(pixel_values=inputs.pixel_values, 
-                                   max_length=100,
+                                   max_length=50,
                                    num_beams=5,
                                    do_sample=True,
-                                   temperature=1.,
-                                   top_k=50,
-                                   top_p=0.6,
-                                   no_repeat_ngram_size=3,
-                                   num_return_sequences=3)
+                                   temperature=2.,
+                                   top_k = 20,
+                                   no_repeat_ngram_size=5,
+                                   num_return_sequences=num_seq)
 
     if tokenizer is not None:
-        generated_caption = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        generated_caption = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
     else:
-        generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)
     return generated_caption
 
-def generate_captions(image):
-    caption_vitgpt = generate_caption(vitgpt_processor, vitgpt_model, image, vitgpt_tokenizer)
+def generate_captions(image, num_seq):
+    caption_vitgpt = generate_caption(vitgpt_processor, vitgpt_model, image, vitgpt_tokenizer, num_seq)
     return caption_vitgpt
 
+st.title('Generate text to your image')
 uploaded_file = st.file_uploader("Upload your image")
+num_seq = st.slider('Return sequences quantity', 1, 5, 2)
+
 if uploaded_file is not None:
-    image = Image.open(uploaded_file)
-    st.image(image)
-    generated_caption = generate_caption(vitgpt_processor, vitgpt_model, image, vitgpt_tokenizer)
-    st.write(generated_caption)
-    
+    if st.button('Generate!'):
+        col1, col2 = st.columns(2)
+        with col1:
+            image = Image.open(uploaded_file)
+            st.image(image)
+        with col2:
+            generated_caption = generate_caption(vitgpt_processor, vitgpt_model, image, vitgpt_tokenizer)
+            for i in generated_caption:
+                st.write(i)