add en translation

app.py

@@ -45,8 +45,10 @@ with st.spinner('Generating image caption ...'):
 
     caption = predict(image)
     image.close()
-    # st.success(f'{caption}')
-    st.header(f'{caption}')
+    caption_en = translator.translate(caption, src='fr', dest='en')
+    st.header('Prediction:')
+    st.subheader(f'{caption}\n')
+    st.subheader(f'English Translation: {caption_en}\n')
 
 st.sidebar.header("ViT-GPT2 predicts:")
 st.sidebar.write(f"{caption}", '\n')

model.py

@@ -7,6 +7,9 @@ from transformers import ViTFeatureExtractor
 from transformers import GPT2Tokenizer
 from huggingface_hub import hf_hub_download
 
+from googletrans import Translator
+translator = Translator()
+
 current_path = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(current_path)
 
@@ -57,6 +60,8 @@ def predict(image):
     while '  ' in caption:
         caption = caption.replace('  ', ' ')
     caption = caption.strip()
+    if caption:
+        caption = caption[0].upper() + caption[1:]
 
     return caption
 

requirements.txt

@@ -3,4 +3,5 @@ Pillow
 jax[cpu]
 flax
 transformers
-huggingface_hub
+huggingface_hub
+googletrans==4.0.0-rc1