layer-rep/Image_Captions

1"""2BLIP Image Captioning  –  Windows + Python 3.12 friendly3--------------------------------------------------------4• Upload an image   OR   drop a webpage URL → get AI captions.5• Model:  Salesforce/blip-image-captioning-base   (no tokenizer crashes)6• Extras: Hero header, copy‑to‑clipboard button, two‑column layout.7"""8 9import gradio as gr10from transformers import BlipProcessor, BlipForConditionalGeneration11from PIL import Image12from io import BytesIO13from bs4 import BeautifulSoup14import requests, torch, os15 16# -------------------------  Load model  ---------------------------------17MODEL_ID = "Salesforce/blip-image-captioning-base"18 19processor = BlipProcessor.from_pretrained(MODEL_ID)20model = BlipForConditionalGeneration.from_pretrained(21    MODEL_ID,22    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float3223)24device = "cuda" if torch.cuda.is_available() else "cpu"25model.to(device).eval()26 27# -------------------------  Core helpers  -------------------------------28def blip_caption(pil_img: Image.Image) -> str:29    """Return a caption for a single PIL image."""30    inputs = processor(images=pil_img.convert("RGB"), return_tensors="pt").to(device)31    with torch.no_grad():32        out = model.generate(**inputs, max_new_tokens=50)33    return processor.decode(out[0], skip_special_tokens=True)34 35def caption_from_url(url: str) -> str:36    """Grab first 3 real images from a webpage and caption them."""37    try:38        soup = BeautifulSoup(requests.get(url, timeout=5).text, "html.parser")39    except Exception as e:40        return f"Failed to fetch page: {e}"41 42    captions, seen = [], set()43    for img in soup.find_all("img"):44        src = img.get("src") or ""45        # resolve relative URLs46        if src.startswith("//"):47            src = "https:" + src48        if not src.startswith(("http://", "https://")):49            continue50        if any(x in src for x in ("svg", "base64", "blank", "1x1")) or src in seen:51            continue52        seen.add(src)53        try:54            raw = Image.open(BytesIO(requests.get(src, timeout=5).content))55            captions.append(f"{src}\n{blip_caption(raw)}")56            if len(captions) == 3:57                break58        except Exception:59            continue60    return "\n\n".join(captions) if captions else "No suitable images found."61 62# -------------------------  UI / Gradio  --------------------------------63CSS = """64.gr-box {border-radius:14px !important;}65button {font-weight:600;}66"""67 68with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:69    # Hero banner70    gr.Markdown("""71    <div style='display:flex;align-items:center;gap:10px;font-size:1.5rem'>72        <span style='font-size:2rem'>🖼️</span><b>BLIP Image Captioning</b>73    </div>74    <p style='margin-top:-8px'>75        Upload an image or paste a webpage URL – I’ll describe what I see.76        <a href='https://huggingface.co/Salesforce/blip-image-captioning-base' target='_blank'>Model card ↗</a>77    </p>78    """)79 80    with gr.Row():81        # ------------  Image column ------------82        with gr.Column(scale=1):83            img_in   = gr.Image(label="Drop an image", type="pil", height=280)84            cap_out  = gr.Textbox(label="Caption", interactive=False)85            with gr.Row():86                btn_cap   = gr.Button("Generate caption")87                copy_btn  = gr.Button("📋 Copy", size="sm", variant="secondary")88            btn_cap.click(lambda i: blip_caption(i) if i else "No image.",89                          img_in, cap_out)90            copy_btn.click(None,        # no Python fn91                           cap_out, []) # no Python outputs92            # copy to clipboard in the browser93            copy_btn.click(js="navigator.clipboard.writeText(args[0]);",  # <-- param is js94                           inputs=cap_out, outputs=None)95 96        # ------------  URL column --------------97        with gr.Column(scale=1):98            url_in  = gr.Textbox(label="Webpage URL")99            url_out = gr.Textbox(label="Captions (first 3 images)", lines=8, interactive=False)100            btn_url = gr.Button("Scrape & caption")101            btn_url.click(caption_from_url, url_in, url_out)102 103demo.launch(share=True)104