OutfitChanger / img2txt.py
altayavci's picture
Upload 17 files
d3fbdbe
raw
history blame
498 Bytes
from transformers import pipeline
captioner = None
PROMPT = "The main subject of this picture is a"
def init():
global captioner
captioner = pipeline(
"image-to-text",
model="Salesforce/blip-image-captioning-base",
prompt=PROMPT
)
def derive_caption(image):
result = captioner(image, max_new_tokens=20)
raw_caption = result[0]["generated_text"]
caption = raw_caption.lower().replace(PROMPT.lower(), "").strip()
return caption