File size: 972 Bytes
c95d497 e828fe3 c95d497 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
import logging
import os
import gradio as gr
from pillow_heif import register_heif_opener
register_heif_opener()
import gradio as gr
from transformers import pipeline
LOG_LEVEL = os.getenv("LOG_LEVEL", "DEBUG")
MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", 200))
# https://huggingface.co/models?pipeline_tag=image-to-text&sort=likes
MODEL = os.getenv("MODEL", "Salesforce/blip-image-captioning-large")
logging.basicConfig(level=LOG_LEVEL)
logger = logging.getLogger(__name__)
logger.info("Loading model...")
# simpler model: "ydshieh/vit-gpt2-coco-en"
captioner = pipeline(
"image-to-text",
model=MODEL,
max_new_tokens=MAX_NEW_TOKENS,
)
logger.info("Done loading model.")
def graptioner(image_url):
global captioner
result = captioner(image_url)
caption = result[0]["generated_text"]
return caption
# add gradio interface
iface = gr.Interface(fn=graptioner, inputs="text", outputs=["text"], allow_flagging="never")
iface.launch()
|