Update pages/ImageToText.py

pages/ImageToText.py

@@ -14,7 +14,15 @@ vitgpt_model.to(device)
 def generate_caption(processor, model, image, tokenizer=None):
     
     inputs = processor(images=image, return_tensors="pt").to(device)
-    generated_ids = model.generate(pixel_values=inputs.pixel_values, max_length=50)
+    generated_ids = model.generate(pixel_values=inputs.pixel_values, 
+                                   max_length=100,
+                                   num_beams=5,
+                                   do_sample=True,
+                                   temperature=1.,
+                                   top_k=50,
+                                   top_p=0.6,
+                                   no_repeat_ngram_size=3,
+                                   num_return_sequences=3)
 
     if tokenizer is not None:
         generated_caption = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
@@ -29,6 +37,7 @@ def generate_captions(image):
 uploaded_file = st.file_uploader("Upload your image")
 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)