Update app.py

app.py

@@ -7,10 +7,17 @@ import torch.nn.functional as F
 from huggingface_hub import hf_hub_download
 from torch.autograd import Variable
 from PIL import Image
+from transformers import pipeline
 
+captioner = pipeline(
+    "image-to-text",
+    model="Salesforce/blip-image-captioning-base",
+    prompt="The main subject of this picture is a"
+)
 
 def removeBackground(image):
     image = Image.fromarray(image).convert("RGB")
+    caption = captioner(image, max_new_tokens=20)[0]["generated_text"].lower().replace("The main subject of this picture is a".lower(), "").strip()
     if not os.path.exists("saved_models"):
         os.mkdir("saved_models")
         os.mkdir("git")
@@ -139,14 +146,10 @@ def removeBackground(image):
     cropped = im_rgb.copy()
     cropped.putalpha(mask)
 
-    return cropped
-
-def remove_background(image):
-    return removeBackground(image)
-
+    return cropped, caption
 
 inputs = gr.inputs.Image()
-outputs = gr.outputs.Image(type="pil")
+outputs = [gr.outputs.Image(type="pil"), "text"]
 interface = gr.Interface(
     fn=remove_background,
     inputs=inputs,