File size: 498 Bytes
d3fbdbe
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
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