|
import os |
|
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False") |
|
import logging |
|
import gradio as gr |
|
from pillow_heif import register_heif_opener |
|
|
|
register_heif_opener() |
|
|
|
from transformers import pipeline |
|
|
|
|
|
LOG_LEVEL = os.getenv("LOG_LEVEL", "DEBUG") |
|
MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", 200)) |
|
|
|
MODEL = os.getenv("MODEL", "Salesforce/blip-image-captioning-large") |
|
|
|
logging.basicConfig(level=LOG_LEVEL) |
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
logger.info("Loading model...") |
|
|
|
captioner = pipeline( |
|
"image-to-text", |
|
model=MODEL, |
|
max_new_tokens=MAX_NEW_TOKENS, |
|
) |
|
logger.info("Done loading model.") |
|
|
|
|
|
|
|
def graptioner(img): |
|
result = captioner(img) |
|
caption = result[0]["generated_text"] |
|
return caption |
|
|
|
iface = gr.Interface( |
|
fn=graptioner, |
|
inputs=gr.components.Image(type="pil"), |
|
outputs=["text"], |
|
allow_flagging="never", |
|
|
|
) |
|
|
|
iface.launch() |
|
|
