CoolFace
Apppublic

Shopify/background-replacement

sourceHugging Faceupdated 2y agoView on Hugging Face
457likes
captioner.py24 linesDownload Raw Back to root
1from transformers import pipeline2 3captioner = None4PROMPT = "The main subject of this picture is a"5 6 7def init():8    global captioner9 10    print("Initializing captioner...")11 12    captioner = pipeline(13        "image-to-text",14        model="Salesforce/blip-image-captioning-base",15        prompt=PROMPT16    )17 18 19def derive_caption(image):20    result = captioner(image, max_new_tokens=20)21    raw_caption = result[0]["generated_text"]22    caption = raw_caption.lower().replace(PROMPT.lower(), "").strip()23    return caption24