Spaces:
Running on Zero
Running on Zero
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import spaces # noqa: E402 must precede torch / transformers | |
| import json | |
| import re | |
| import html | |
| import tempfile | |
| import torch | |
| import gradio as gr | |
| from PIL import Image, ImageDraw, ImageFont | |
| from transformers import AutoProcessor, AutoModelForImageTextToText | |
| MODEL_ID = "P1n3/sdg-detector-grpo" | |
| # --------------------------------------------------------------------------- | |
| # Model (loaded once at module scope, moved to CUDA eagerly for ZeroGPU) | |
| # --------------------------------------------------------------------------- | |
| processor = AutoProcessor.from_pretrained(MODEL_ID) | |
| model = AutoModelForImageTextToText.from_pretrained( | |
| MODEL_ID, | |
| dtype=torch.bfloat16, | |
| attn_implementation="sdpa", | |
| ).to("cuda") | |
| model.eval() | |
| # --------------------------------------------------------------------------- | |
| # Prompt (the SDG detector's structured-grounding question template) | |
| # --------------------------------------------------------------------------- | |
| QUESTION_TEMPLATE = """You are an AI image quality evaluator. You will be given **one image** to analyze. | |
| ### Definitions | |
| **Misalignment**: Areas where the image content does NOT match the text caption, including: | |
| - Missing objects: Objects mentioned in caption but not present in image | |
| - Extra objects: Objects present in image but not mentioned in caption | |
| - Wrong attributes: Incorrect color, size, material, count, or other properties | |
| - Wrong spatial relationships: Incorrect positions, orientations, or arrangements | |
| **Artifact**: Visual defects in the generated image, including: | |
| - Distorted anatomy: Malformed hands, extra/missing limbs, wrong number of fingers | |
| - Duplicated/missing parts: Repeated or absent body parts, objects | |
| - Warped geometry: Perspective errors, impossible shapes | |
| - Texture issues: Melted, smeared, or overly smooth textures | |
| - Unnatural edges: Jagged, broken, or blurry boundaries | |
| - Garbled text: Unreadable or malformed text/letters | |
| - Lighting inconsistencies: Wrong shadows, reflections, or light sources | |
| Text Caption: {caption} | |
| **Goal**: Produce a detailed analysis of the image quality and output bounding boxes for all detected issues. | |
| ### Strict Output Rules | |
| Output **TWO blocks in this exact order**: | |
| 1) `<think>` - Your detailed analysis | |
| 2) `<answer>` - JSON list of bounding boxes | |
| ### Answer Format (for <answer>) | |
| Return a JSON list: | |
| [ | |
| {{"box_2d": [x0, y0, x1, y1], "label": "misalignment"|"artifact", "description": "brief description of the issue", "importance": 1-100}} | |
| ] | |
| Bounding box coordinates are in normalized 0-1000 space: [x0, y0, x1, y1]. | |
| The "importance" is an integer from 1 (minor) to 100 (severe) rating how much the defect hurts image quality. | |
| If there are no issues, output an empty list. | |
| Now analyze the image and produce your output: | |
| """ | |
| ARTIFACT_COLOR = (239, 68, 68) # red | |
| MISALIGNMENT_COLOR = (37, 99, 235) # blue | |
| def _load_font(size): | |
| for path in ( | |
| "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", | |
| "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", | |
| ): | |
| try: | |
| return ImageFont.truetype(path, size) | |
| except Exception: | |
| continue | |
| return ImageFont.load_default() | |
| def parse_answer(response: str): | |
| """Extract the structured defect list from the model's <answer> block.""" | |
| m = re.search(r"<answer>\s*(.*?)\s*</answer>", response, re.DOTALL) | |
| payload = m.group(1) if m else response | |
| start, end = payload.find("["), payload.rfind("]") | |
| if start == -1 or end == -1 or end <= start: | |
| return [] | |
| try: | |
| data = json.loads(payload[start:end + 1]) | |
| except Exception: | |
| return [] | |
| if not isinstance(data, list): | |
| return [] | |
| defects = [] | |
| for item in data: | |
| if not isinstance(item, dict): | |
| continue | |
| box = item.get("box_2d") | |
| if not (isinstance(box, list) and len(box) == 4): | |
| continue | |
| try: | |
| box = [float(v) for v in box] | |
| except Exception: | |
| continue | |
| label = str(item.get("label", "")).lower() | |
| if label not in ("artifact", "misalignment"): | |
| label = "artifact" | |
| desc = item.get("description") or item.get("desc") or "" | |
| imp = item.get("importance") | |
| try: | |
| imp = int(imp) | |
| except Exception: | |
| imp = None | |
| defects.append({ | |
| "box_2d": box, | |
| "label": label, | |
| "description": str(desc), | |
| "importance": imp, | |
| }) | |
| return defects | |
| def extract_think(response: str) -> str: | |
| m = re.search(r"<think>\s*(.*?)\s*</think>", response, re.DOTALL) | |
| return m.group(1).strip() if m else "" | |
| def draw_defects(image: Image.Image, defects): | |
| """Overlay defect bounding boxes (coords normalized 0-1000, Qwen order).""" | |
| image = image.convert("RGB").copy() | |
| w, h = image.size | |
| draw = ImageDraw.Draw(image) | |
| line_w = max(3, round(min(w, h) / 250)) | |
| font = _load_font(max(16, round(min(w, h) / 45))) | |
| for idx, d in enumerate(defects, start=1): | |
| x0, y0, x1, y1 = d["box_2d"] | |
| px0, py0 = x0 / 1000 * w, y0 / 1000 * h | |
| px1, py1 = x1 / 1000 * w, y1 / 1000 * h | |
| if px1 < px0: | |
| px0, px1 = px1, px0 | |
| if py1 < py0: | |
| py0, py1 = py1, py0 | |
| color = ARTIFACT_COLOR if d["label"] == "artifact" else MISALIGNMENT_COLOR | |
| draw.rectangle([px0, py0, px1, py1], outline=color, width=line_w) | |
| imp = d["importance"] | |
| tag = f"{idx}" + (f" \u00b7 {imp}" if imp is not None else "") | |
| bbox = draw.textbbox((0, 0), tag, font=font) | |
| tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1] | |
| ly = max(0, py0 - th - 4) | |
| draw.rectangle([px0, ly, px0 + tw + 8, ly + th + 4], fill=color) | |
| draw.text((px0 + 4, ly + 2), tag, fill=(255, 255, 255), font=font) | |
| return image | |
| def defects_to_markdown(defects): | |
| if not defects: | |
| return ("### ✅ No defects detected\n\n" | |
| "The SDG detector did not ground any localized artifacts or " | |
| "caption misalignments in this image.") | |
| n_art = sum(1 for d in defects if d["label"] == "artifact") | |
| n_mis = sum(1 for d in defects if d["label"] == "misalignment") | |
| lines = [ | |
| f"### Found {len(defects)} defect(s) — " | |
| f"🔴 {n_art} artifact, 🔵 {n_mis} misalignment\n", | |
| "| # | Type | Importance | Where / What / Why |", | |
| "|---|------|-----------|--------------------|", | |
| ] | |
| for idx, d in enumerate(defects, start=1): | |
| emoji = "🔴" if d["label"] == "artifact" else "🔵" | |
| imp = d["importance"] if d["importance"] is not None else "—" | |
| box = ", ".join(str(int(v)) for v in d["box_2d"]) | |
| desc = html.escape(d["description"]).replace("\n", " ").replace("|", "\\|") | |
| lines.append( | |
| f"| {idx} | {emoji} {d['label']} | {imp} | " | |
| f"{desc} <br><sub>box (0–1000): [{box}]</sub> |" | |
| ) | |
| return "\n".join(lines) | |
| def detect(image, caption, max_new_tokens=1024): | |
| """Detect and ground structured defects in a text-to-image generation. | |
| Args: | |
| image: the generated image to inspect for defects. | |
| caption: the text prompt the image was generated from (used to spot | |
| caption/image misalignments). Optional. | |
| max_new_tokens: generation budget for the detector's reasoning + answer. | |
| Returns: | |
| The image annotated with defect boxes, a structured defect table, and | |
| the detector's raw reasoning. | |
| """ | |
| if image is None: | |
| raise gr.Error("Please provide an image to analyze.") | |
| caption = (caption or "").strip() | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "image", "image": image}, | |
| {"type": "text", "text": QUESTION_TEMPLATE.format(caption=caption)}, | |
| ], | |
| } | |
| ] | |
| inputs = processor.apply_chat_template( | |
| messages, | |
| tokenize=True, | |
| add_generation_prompt=True, | |
| return_dict=True, | |
| return_tensors="pt", | |
| ).to(model.device) | |
| with torch.inference_mode(): | |
| generated = model.generate( | |
| **inputs, | |
| max_new_tokens=int(max_new_tokens), | |
| do_sample=False, | |
| ) | |
| trimmed = generated[:, inputs["input_ids"].shape[1]:] | |
| response = processor.batch_decode( | |
| trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False | |
| )[0] | |
| defects = parse_answer(response) | |
| annotated = draw_defects(image, defects) | |
| table = defects_to_markdown(defects) | |
| think = extract_think(response) | |
| reasoning = think if think else response.strip() | |
| return annotated, table, reasoning | |
| CSS = """ | |
| #col-container { max-width: 1200px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| INTRO = """# 🔍 Structured Defect Grounding (SDG) | |
| Detect **localized defects** in text-to-image generations and get instance-level | |
| feedback — *where* the defect is, *what* type it is (🔴 artifact / 🔵 misalignment), | |
| *why* it's wrong, and *how important* it is (1–100). | |
| Powered by the **SDG Detector** ([`P1n3/sdg-detector-grpo`](https://huggingface.co/P1n3/sdg-detector-grpo)), | |
| a Qwen3-VL-4B model trained with SFT + GRPO from the paper | |
| [*Where, What, Why, and Importance: Structured Defect Grounding for Text-to-Image Feedback*](https://huggingface.co/papers/2606.06113). | |
| Upload a generated image and (optionally) the caption it was generated from, | |
| then press **Analyze**. | |
| """ | |
| with gr.Blocks(title="Structured Defect Grounding") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown(INTRO) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| image_in = gr.Image(type="pil", label="Generated image", height=380) | |
| caption_in = gr.Textbox( | |
| label="Caption / prompt (optional)", | |
| placeholder="The text prompt the image was generated from…", | |
| lines=2, | |
| ) | |
| run = gr.Button("Analyze", variant="primary") | |
| with gr.Accordion("Advanced settings", open=False): | |
| max_tokens = gr.Slider( | |
| 256, 2048, value=1024, step=64, | |
| label="Max new tokens", | |
| info="Generation budget for the detector's reasoning + answer.", | |
| ) | |
| with gr.Column(scale=1): | |
| image_out = gr.Image(type="pil", label="Detected defects", height=380) | |
| table_out = gr.Markdown(label="Structured feedback") | |
| with gr.Accordion("Detector reasoning (<think>)", open=False): | |
| reasoning_out = gr.Textbox(label="Raw reasoning", lines=8) | |
| gr.Examples( | |
| examples=[ | |
| ["examples/throne_dystopia.png", "a jung male sitting down on a throne in a dystopian world, digital art, epic"], | |
| ["examples/sign_bianca_buda.png", "A sign that says Bianca Buda"], | |
| ["examples/pikachu_mario.png", "Pikachu Ninja turtle Mewtwo super Mario"], | |
| ["examples/car_fruit_stand.png", "a car driving through a fruit stand, movie action scene, fruits flying everywhere"], | |
| ], | |
| inputs=[image_in, caption_in], | |
| outputs=[image_out, table_out, reasoning_out], | |
| fn=detect, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| run.click( | |
| detect, | |
| inputs=[image_in, caption_in, max_tokens], | |
| outputs=[image_out, table_out, reasoning_out], | |
| api_name="detect", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |