Update app.py

app.py

@@ -7,24 +7,35 @@ model, _, transform = open_clip.create_model_and_transforms(
   model_name="coca_ViT-L-14",
   pretrained="laion2B-s13B-b90k"
 )
-
+def resize_background(img):
+    width, height = img.size
+    aspect_ratio = width / height
+    
+    if aspect_ratio != 1:
+        new_img = Image.new("RGB", (512, 512), color=(255, 255, 255))
+        new_img.paste(img, (int((512 - width) / 2), int((512 - height) / 2)))
+        img = new_img
+    
+    img = img.resize((512, 512), Image.LANCZOS)
+    return img
 # load an image
 def generate_caption(image):
-    im = image.convert("RGB")
+    im = resize_background(image)
     # transform the image and add a batch size dimension
     im = transform(im).unsqueeze(0)
 
     generated = model.generate(im)
     generated = generated.detach()
 
-    return(open_clip.decode(generated[0]))
+    return(open_clip.decode(generated[0]).split("<start_of_text>")[1].split("<end_of_text>")[0])
 
 with gr.Blocks() as demo:
     gr.Markdown("## Captioning with OpenCLIP CoCa")
     with gr.Row():
-        inp = gr.Image(label="Image to Caption", type="pil")
+        with gr.Column():
+            image = gr.Image(label="Image to Caption", type="pil")
         out = gr.Textbox(label="Caption")
     btn = gr.Button("Generate caption")
-    btn.click(fn=generate_caption,inputs=inp, outputs=out)
+    btn.click(fn=generate_caption,inputs=image, outputs=out)
 
 demo.launch()