|
import argparse |
|
import io |
|
import logging |
|
import os |
|
|
|
import gradio as gr |
|
import requests |
|
from PIL import Image |
|
from pillow_heif import register_heif_opener |
|
from transformers import pipeline |
|
|
|
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False") |
|
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") |
|
|
|
register_heif_opener() |
|
|
|
logging.basicConfig(level=LOG_LEVEL) |
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
def setup_args(): |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument("--share", action="store_true", default=False) |
|
return parser.parse_args() |
|
|
|
|
|
def load_image_from_url(url): |
|
try: |
|
response = requests.get(url) |
|
if not response.ok: |
|
raise Exception("Error downloading image") |
|
image = Image.open(io.BytesIO(response.content)) |
|
return image |
|
except Exception as e: |
|
logger.error("Error loading image from URL: %s", e) |
|
raise |
|
|
|
|
|
def graptioner(image, url): |
|
if url and url.strip(): |
|
image = load_image_from_url(url) |
|
width, height = image.size |
|
if width < 1 or height < 1: |
|
raise Exception("Invalid image") |
|
logger.debug("Loaded image size: %sx%s", width, height) |
|
|
|
result = captioner(image) |
|
return result[0]["generated_text"] |
|
|
|
|
|
if __name__ == "__main__": |
|
args = setup_args() |
|
logger.info("Loading model...") |
|
|
|
captioner = pipeline( |
|
"image-to-text", |
|
model=MODEL, |
|
max_new_tokens=MAX_NEW_TOKENS, |
|
) |
|
logger.info("Done loading model.") |
|
iface = gr.Interface( |
|
fn=graptioner, |
|
inputs=[ |
|
gr.Image(type="pil", label="Upload Image"), |
|
gr.Textbox(lines=1, placeholder="Image URL", label="Image URL"), |
|
], |
|
outputs=["text"], |
|
allow_flagging="never", |
|
) |
|
iface.launch(share=args.share) |
|
|
