ghuflabs/Vision_Language_Studio

1import gradio as gr2import torch3import numpy as np4from PIL import Image, ImageDraw, ImageOps, ImageFilter5 6# Optional: enable ZeroGPU with @spaces.GPU if you want7# import spaces8 9DEVICE = "cuda" if torch.cuda.is_available() else "cpu"10DTYPE = torch.float16 if torch.cuda.is_available() else torch.float3211 12_models = {}13 14def get_captioner():15    if "captioner" not in _models:16        from transformers import BlipProcessor, BlipForConditionalGeneration17        name = "Salesforce/blip-image-captioning-base"18        processor = BlipProcessor.from_pretrained(name)19        model = BlipForConditionalGeneration.from_pretrained(name, torch_dtype=DTYPE).to(DEVICE)20        _models["captioner"] = (processor, model)21    return _models["captioner"]22 23def get_vqa():24    if "vqa" not in _models:25        from transformers import ViltProcessor, ViltForQuestionAnswering26        name = "dandelin/vilt-b32-finetuned-vqa"27        processor = ViltProcessor.from_pretrained(name)28        model = ViltForQuestionAnswering.from_pretrained(name, torch_dtype=DTYPE).to(DEVICE)29        _models["vqa"] = (processor, model)30    return _models["vqa"]31 32def get_ocr():33    if "ocr" not in _models:34        from transformers import TrOCRProcessor, VisionEncoderDecoderModel35        name = "microsoft/trocr-base-printed"36        processor = TrOCRProcessor.from_pretrained(name)37        model = VisionEncoderDecoderModel.from_pretrained(name, torch_dtype=DTYPE).to(DEVICE)38        _models["ocr"] = (processor, model)39    return _models["ocr"]40 41def get_clip():42    if "clip" not in _models:43        from transformers import CLIPProcessor, CLIPModel44        name = "openai/clip-vit-base-patch32"45        processor = CLIPProcessor.from_pretrained(name)46        model = CLIPModel.from_pretrained(name, torch_dtype=DTYPE).to(DEVICE)47        _models["clip"] = (processor, model)48    return _models["clip"]49 50MAX_IMAGE_SIZE = 102451MAX_TOKENS = 6452 53# ---------- helpers ----------54def _parse_candidates(txt: str):55    if not txt:56        return []57    return [ln.strip() for ln in txt.splitlines() if ln.strip()]58 59def _prep_ocr_image(img: Image.Image) -> Image.Image:60    # Robust printed-text normalization for TrOCR61    g = img.convert("L")62    g = ImageOps.autocontrast(g)63    # light denoise + sharpen64    g = g.filter(ImageFilter.MedianFilter(size=3))65    g = ImageOps.invert(ImageOps.invert(g).filter(ImageFilter.UnsharpMask(radius=1.2, percent=130)))66    # light binarization while keeping grayscale nuance67    arr = np.array(g, dtype=np.uint8)68    thr = int(np.clip(arr.mean() * 0.85, 100, 200))69    arr = np.where(arr > thr, 255, arr)70    g2 = Image.fromarray(arr.astype(np.uint8), mode="L").convert("RGB")71    # resize toward TrOCR’s typical scale72    longest = 38473    w, h = g2.size74    scale = longest / max(w, h)75    if scale < 1.0:76        g2 = g2.resize((int(w * scale), int(h * scale)), Image.BICUBIC)77    return g278 79# --- Inference functions ---80 81def infer_caption(img, max_new_tokens, num_beams, progress=gr.Progress(track_tqdm=True)):82    if img is None:83        return "Please upload an image."84    processor, model = get_captioner()85    inputs = processor(images=img, return_tensors="pt").to(DEVICE)86    with torch.no_grad():87        out = model.generate(**inputs, max_new_tokens=int(max_new_tokens), num_beams=int(num_beams))88    caption = processor.decode(out[0], skip_special_tokens=True)89    return caption90 91def infer_vqa(img, question, progress=gr.Progress(track_tqdm=True)):92    if img is None:93        return "Please upload an image."94    if not question or not question.strip():95        return "Please type a question."96    processor, model = get_vqa()97    enc = processor(img, question, return_tensors="pt").to(DEVICE)98    with torch.no_grad():99        logits = model(**enc).logits100    idx = int(logits.argmax(-1))101    return model.config.id2label[idx]102 103def infer_ocr(img, progress=gr.Progress(track_tqdm=True)):104    if img is None:105        return "Please upload an image."106    img_proc = _prep_ocr_image(img)107    processor, model = get_ocr()108    pixel_values = processor(images=img_proc, return_tensors="pt").pixel_values.to(DEVICE)109    with torch.no_grad():110        ids = model.generate(111            pixel_values,112            max_new_tokens=MAX_TOKENS,113            do_sample=True,114            num_beams=3,115            temperature=1.0,116            repetition_penalty=1.05,117        )118    text = processor.batch_decode(ids, skip_special_tokens=True)[0]119    text = text.strip()120 121    if len(text) > 2:122        upper_ratio = sum(1 for c in text if c.isupper()) / max(1, sum(c.isalpha() for c in text))123        if upper_ratio > 0.8:  124            text = text.lower()125            import re126            text = re.sub(r"(^[a-z])", lambda m: m.group(1).upper(), text)127            text = re.sub(r"([.!?]\s*)([a-z])", lambda m: m.group(1) + m.group(2).upper(), text)128    return text if text else "(no text detected)"129 130def infer_clip_match(img, candidate_texts, topk=5, progress=gr.Progress(track_tqdm=True)):131    if img is None:132        return "Please upload an image.", ""133    candidates = _parse_candidates(candidate_texts)134    if not candidates:135        return "Please provide candidate texts (one per line).", ""136 137    processor, model = get_clip()138 139    try:140        with torch.no_grad():141            image_inputs = processor(images=img, return_tensors="pt").to(DEVICE)142            image_features = model.get_image_features(**image_inputs)143            image_features = image_features / image_features.norm(p=2, dim=-1, keepdim=True)144    except Exception as e:145        return f"Failed to encode image with CLIP: {e}", ""146 147    all_text_features = []148    chunk_size = 24149    try:150        for i in range(0, len(candidates), chunk_size):151            chunk = candidates[i:i+chunk_size]152            with torch.no_grad():153                text_inputs = processor(text=chunk, return_tensors="pt", padding=True, truncation=True).to(DEVICE)154                text_features = model.get_text_features(**text_inputs)155                text_features = text_features / text_features.norm(p=2, dim=-1, keepdim=True)156            all_text_features.append(text_features)157        text_features = torch.cat(all_text_features, dim=0)158    except Exception as e:159        return f"Failed to encode text with CLIP: {e}", ""160 161    try:162        sims = (image_features @ text_features.T).squeeze(0)  # (num_text,)163        sims_cpu = sims.float().cpu().numpy()164        order = sims_cpu.argsort()[::-1]165        k = int(min(max(1, int(topk)), len(order)))166        top_idx = order[:k]167        md_lines = ["### Ranking"]168        for rank, idx in enumerate(top_idx, start=1):169            md_lines.append(f"{rank}. **{candidates[int(idx)]}** — score: `{sims_cpu[int(idx)]:.3f}`")170        best_text = candidates[int(top_idx[0])]171        md = "\n".join(md_lines)172        return md, best_text173    except Exception as e:174        return f"Failed to compute similarity: {e}", ""175 176# ---------- UI ----------177css = """178#col { margin: 0 auto; max-width: 1080px; }179"""180 181with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:182    with gr.Column(elem_id="col"):183        gr.Markdown("# Vision Language Studio\nImage captioning, VQA, OCR, and CLIP matching")184 185        with gr.Tab("Image Captioning"):186            cap_img = gr.Image(type="pil", label="Image")187            cap_tokens = gr.Slider(label="Max new tokens", minimum=8, maximum=64, step=1, value=30)188            cap_beams = gr.Slider(label="Beam width", minimum=1, maximum=5, step=1, value=3)189            cap_btn = gr.Button("Generate Caption", variant="primary")190            cap_out = gr.Textbox(label="Caption", max_lines=2, show_label=True)191            gr.on([cap_btn.click], infer_caption, [cap_img, cap_tokens, cap_beams], [cap_out])192 193        with gr.Tab("Visual Q&A"):194            with gr.Row():195                vqa_img = gr.Image(type="pil", label="Image")196                with gr.Column(scale=1):197                    vqa_q = gr.Textbox(label="Question", placeholder="What is in the image?")198                    vqa_btn = gr.Button("Answer", variant="primary")199                    vqa_out = gr.Textbox(label="Answer", max_lines=2)200            with gr.Row():201                vqa_example = gr.Dropdown(202                    label="Load example",203                    choices=[204                        "sample_images/cat.png | What animal is this?",205                        "sample_images/bike.png | Is this person riding a bike or a motorcycle?",206                    ],207                    value=None,208                    interactive=True,209                )210 211            def _load_vqa_example(sel):212                if not sel:213                    return None, ""214                path, q = [x.strip() for x in sel.split("|")]215                try:216                    img = Image.open(path).convert("RGB")217                except Exception:218                    img = None219                return img, q220 221            vqa_example.change(_load_vqa_example, vqa_example, [vqa_img, vqa_q])222            gr.on([vqa_btn.click, vqa_q.submit], infer_vqa, [vqa_img, vqa_q], [vqa_out])223 224        with gr.Tab("OCR (Printed Text)"):225            with gr.Row():226                ocr_img = gr.Image(type="pil", label="Image with printed text")227                with gr.Column(scale=1):228                    ocr_btn = gr.Button("Extract Text", variant="primary")229                    ocr_out = gr.Textbox(label="Recognized text", max_lines=6)230            gr.on([ocr_btn.click], infer_ocr, [ocr_img], [ocr_out])231 232        with gr.Tab("CLIP VLM"):233            clip_img = gr.Image(type="pil", label="Image")234            clip_candidates = gr.Textbox(235                label="Candidate texts (one per line)",236                value="a photo of a cat\na photo of a dog\nindoor scene\noutdoor landscape\na man riding a bike",237                lines=8,238            )239            clip_topk = gr.Slider(label="Show top-k", minimum=1, maximum=10, step=1, value=5)240            clip_btn = gr.Button("Match Image ↔ Text", variant="primary")241            clip_md = gr.Markdown()242            clip_best = gr.Textbox(label="Best match", max_lines=1)243 244            gr.on(245                [clip_btn.click],246                infer_clip_match,247                [clip_img, clip_candidates, clip_topk],248                [clip_md, clip_best]249            )250 251if __name__ == "__main__":252    demo.launch()253