Update app.py

app.py

@@ -43,32 +43,53 @@ model_vq = blip_vqa(pretrained=model_url_vq, image_size=480, vit='base')
 model_vq.eval()
 model_vq = model_vq.to(device)
 
-
-
-def inference(raw_image, model_n, question, strategy):
+def getModelPath(language):
+    if language == 'English':
+        path = None
+    elif language == 'German':
+        path = "Helsinki-NLP/opus-mt-en-de"
+    elif language == 'French':
+        path = "Helsinki-NLP/opus-mt-en-fr"
+    elif language == 'Spanish':
+        path = "Helsinki-NLP/opus-mt-en-es"
+    elif language == 'Chinese': 
+        path = "Helsinki-NLP/opus-mt-en-zh"
+    elif language == 'Ukranian':
+        path = "Helsinki-NLP/opus-mt-en-uk"
+    elif language == 'Swedish':
+        path = "Helsinki-NLP/opus-mt-en-sv"
+    elif language == 'Arabic':
+        path = "Helsinki-NLP/opus-mt-en-ar"
+    elif language == 'Italian':
+        path = "Helsinki-NLP/opus-mt-en-it"
+    elif language == 'Hindi':
+        path = "Helsinki-NLP/opus-mt-en-hi"
+    return(path)
+
+def inference(input_img,strategy,language):
     if model_n == 'Image Captioning':
         image = transform(raw_image).unsqueeze(0).to(device)   
         with torch.no_grad():
-          if strategy == "Beam search":
-              caption = model.generate(image, sample=False, num_beams=3, max_length=20, min_length=5)
-          else:
-              caption = model.generate(image, sample=True, top_p=0.9, max_length=20, min_length=5)
-          return 'caption: '+caption[0]
-
-    else:   
-        image_vq = transform_vq(raw_image).unsqueeze(0).to(device)  
-        with torch.no_grad():
-            answer = model_vq(image_vq, question, train=False, inference='generate') 
-        return  'answer: '+answer[0]
+            if strategy == "Beam search":
+              cap = model.generate(image, sample=False, num_beams=3, max_length=20, min_length=5)
+            else:
+              cap = model.generate(image, sample=True, top_p=0.9, max_length=20, min_length=5)
+            if modelpath:
+                translator = pipeline("translation", model=modelpath)
+                trans_cap = translator(cap[0])
+                tc = trans_cap[0]['translation_text']
+                return str(tc)
+            else:
+                return str(cap[0])
+
     
-inputs = [gr.inputs.Image(type='pil'),gr.inputs.Radio(choices=['Image Captioning',"Visual Question Answering"], type="value", default="Image Captioning", label="Task"),gr.inputs.Textbox(lines=2, label="Question"),gr.inputs.Radio(choices=['Beam search','Nucleus sampling'], type="value", default="Nucleus sampling", label="Caption Decoding Strategy")]
-outputs = gr.outputs.Textbox(label="Output")
+description = "A pipeline of BLIP image captioning and Helsinki translation in order to generate image captions in a language of your choice either with beam search (deterministic) or nucleus sampling (stochastic). Enjoy! Is the language you want to use missing? Let me know and I'll integrate it."
 
-title = "BLIP"
 
-description = "Gradio demo for BLIP: Bootstrapping Language-Image Pre-training for Unified Vision-Language Understanding and Generation (Salesforce Research). To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
+inputs_ = [gr.inputs.Image(type='filepath', label="Input Image"),gr.inputs.Radio(choices=['Beam search','Nucleus sampling'], type="value", default="Nucleus sampling", label="Mode"), gr.inputs.Radio(choices=['English','German', 'French', 'Spanish', 'Chinese', 'Ukranian', 'Swedish', 'Arabic', 'Italian', 'Hindi'],type="value", default = 'German',label="Language")]
 
-article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2201.12086' target='_blank'>BLIP: Bootstrapping Language-Image Pre-training for Unified Vision-Language Understanding and Generation</a> | <a href='https://github.com/salesforce/BLIP' target='_blank'>Github Repo</a></p>"
+outputs_ = gr.outputs.Textbox(label="Output")
 
+iface = gr.Interface(inference, inputs_, outputs_, description=description)
 
-gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, examples=[['starrynight.jpeg',"Image Captioning","None","Nucleus sampling"]]).launch(enable_queue=True)
+iface.launch(debug=True,show_error=True)