code_inference/rule_13_vllm.py

# -*- coding: utf-8 -*-
"""
使用方法:
python rule_13_vllm.py \
  --input_dir "/path/to/your/images" \
  --model_path "/path/to/your/Qwen2.5-VL-7B-Instruct"
"""

import os
import re
import json
import argparse
import multiprocessing
from pathlib import Path
from typing import Dict, List, Optional, Any
from tqdm import tqdm
from PIL import Image, ImageFile
from transformers import AutoProcessor
from vllm import LLM, SamplingParams

os.environ['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn'

ImageFile.LOAD_TRUNCATED_IMAGES = True

IMG_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tif", ".tiff"}

#排布位置
SYS_PROMPT_TEXT="""You are a highly critical Senior Art Director specializing in Layout and Visual Hierarchy. 
Your job is to identify "Suffocating Designs"—creative pieces where elements are too cramped, lack breathing room, or feel disorganized due to poor spacing.

INPUT: One image and one natural-language question about layout composition and spacing.

YOUR TASK:
1. Determine if the image is a **VIOLATION** (Unsuitable) or **SAFE** (Suitable) based on the criteria below.
2. Output a JSON object containing a rigorous Chain-of-Thought ("think") and a precise classification label ("answer").

CORE JUDGMENT PRINCIPLE: A professional advertisement must have a clear "Sense of Breath" (Negative Space). If the elements feel "squeezed" or "crowded," it is UNSUITABLE.

OUTPUT FORMAT:
Return EXACTLY two blocks, no extra text:
<think>Detailed reasoning evaluating negative space, element proximity, visual path, and edge tension...</think><answer>{"Answer": "<Suitable OR Unsuitable>", "Answer type": "Text Legibility and Placement"}</answer>

=========================================
STRICT VIOLATION CRITERIA (If ANY match -> Unsuitable)
=========================================
1. **Lack of Breathing Room (Core Crowding Violation):**
   - **Major Module Conflict:** The main subject (product/hero), the headline text, and the logo are placed too close to each other, lacking deliberate negative space.
   - **Claustrophobic Feel:** The overall design feels "heavy" or "squeezed" because these major elements are fighting for space.

2. **The "Clinging" Small Print (Secondary Text Violation):**
   - **Tangency Risk:** While small text *can* be closer than headlines, it becomes a VIOLATION if it is "clinging" to, touching, or "tangent" to other elements or the image border.
   - **Visual Noise:** Small text is squeezed into gaps without enough margin, looking like an afterthought rather than a design choice.

3. **Edge Tension (贴边风险):**
   - Elements are "touching" the canvas border or each other without a clear, intentional overlap (e.g., accidental contact).

4. **Information Overload (信息堆砌):**
   - **No Visual Path:** The layout is filled with too many text blocks or icons with no clear separation or hierarchy. The eye has nowhere to rest.

5. **Placement & Background Conflict (文字-排布位置与背景):** [NEW CRITICAL RULE]
   - **Text on Noise (背景干扰):** Text is overlaid directly onto a complex, textured, or high-contrast background (e.g., tree branches, detailed patterns) without a drop shadow or mask, making it "hard to breathe/read".
   - **Weak Visual Anchor (视线捕捉失败):** Important text (Headline) is placed in a "dead zone" (extreme edges/corners) or blends into the background, failing to capture the eye immediately.

=========================================
CRITERIA FOR 'SUITABLE' (NON-VIOLATION / GOOD DESIGN)
=========================================
1. **Generous White Space (Major Elements):**
   - Clear and deliberate separation exists between the Headline, Main Subject, and Footer.

2. **Permissible Density (Small Print Exemption):**
   - **Nuance:** Secondary small text (annotations/footnotes) IS ALLOWED to have smaller gaps relative to other elements (unlike headlines). As long as it doesn't touch/cling (see Violation #2), tighter spacing for small text is SAFE.

3. **Valid Exclusions:**
   - **Product Packaging:** Text printed naturally on the product packaging is SAFE.
   - **Artistic Integration:** Artistic fonts visually integrated *into* the product itself are SAFE.
   - **Media Exemption:** Film stills or variety show photography are always SAFE.

=========================================
DECISION LOGIC
=========================================
- **Unsuitable**: If the design feels squeezed, suffers from edge tension, lacks a visual path, or if small text "clings" to edges/elements.
- **Suitable**: If major elements breathe well, OR if the density is strictly limited to allowed small print/packaging text that doesn't create tension.
"""


def collect_images(input_dir: Path) -> List[Dict[str, str]]:
    """扫描目录下所有图片"""
    if not input_dir.exists():
        raise FileNotFoundError(f"Input directory not found: {input_dir}")
    
    files = [p for p in input_dir.iterdir() if p.is_file() and p.suffix.lower() in IMG_EXTS]
    files.sort()
    
    print(f"[Info] Found {len(files)} images in {input_dir}")
    return [{"path": str(p), "filename": p.name} for p in files]

def parse_llm_output(text: str) -> Dict[str, Any]:
    default_res = {
        "label": "Parse Error", 
        "think": "No reasoning found", 
        "raw": text
    }
    
    if not text:
        return default_res

    think_match = re.search(r'<think>(.*?)</think>', text, re.DOTALL)
    think_content = think_match.group(1).strip() if think_match else ""

    answer_match = re.search(r'<answer>(.*?)</answer>', text, re.DOTALL)
    
    extracted_label = "Parse Error"
    
    if answer_match:
        json_str = answer_match.group(1).strip()
        try:

            data = json.loads(json_str)

            raw_ans = data.get("Answer", "")
            
            if "unsuitable" in raw_ans.lower():
                extracted_label = "Unsuitable"
            elif "suitable" in raw_ans.lower():
                extracted_label = "Suitable"
            else:
                extracted_label = raw_ans 
                
        except json.JSONDecodeError:
            if "Unsuitable" in json_str:
                extracted_label = "Unsuitable"
            elif "Suitable" in json_str:
                extracted_label = "Suitable"
    else:
        if "Unsuitable" in text:
            extracted_label = "Unsuitable"
        elif "Suitable" in text:
            extracted_label = "Suitable"

    return {
        "label": extracted_label,
        "think": think_content,
        "raw": text
    }

def prepare_vllm_inputs(batch_meta: List[Dict], processor) -> List[Dict]:
    
    vllm_inputs = []
    user_query = "Analyze this image against the design rules and return the JSON decision."
    
    for item in batch_meta:
        img_path = item["path"]
        try:
            image_obj = Image.open(img_path).convert("RGB")
            
            messages = [
                {"role": "system", "content": [{"type": "text", "text": SYS_PROMPT_TEXT}]},
                {"role": "user", "content": [
                    {"type": "image", "image": img_path},
                    {"type": "text", "text": user_query}
                ]}
            ]
            
            prompt_text = processor.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
            
            vllm_inputs.append({
                "prompt": prompt_text,
                "multi_modal_data": {"image": image_obj}
            })
        except Exception as e:
            print(f"[Warning] Failed to load {img_path}: {e}")
            vllm_inputs.append(None)
            
    return vllm_inputs

# ==========================================
# 主程序
# ==========================================

def main():
    parser = argparse.ArgumentParser(description="AI Visual Comfort Auditor")
    parser.add_argument("--input_dir", type=str, required=True, help="Folder containing images to check")
    parser.add_argument("--model_path", type=str, required=True, help="Path to local Qwen-VL model")
    parser.add_argument("--batch_size", type=int, default=512, help="Inference batch size")
    parser.add_argument("--tp_size", type=int, default=2, help="Tensor Parallel size")
    args = parser.parse_args()

    input_path = Path(args.input_dir)
    meta_data = collect_images(input_path)
    
    if not meta_data:
        print("[Info] No images found. Exiting.")
        return
    print(f"\n[Init] Loading Model: {args.model_path}")
    
    llm = LLM(
        model=args.model_path,
        tokenizer=args.model_path,
        trust_remote_code=True,
        tensor_parallel_size=args.tp_size,
        gpu_memory_utilization=0.90, 
        max_model_len=8192, 
        enforce_eager=True,
        limit_mm_per_prompt={"image": 1}
    )
    
    processor = AutoProcessor.from_pretrained(args.model_path, trust_remote_code=True)
    

    sampling_params = SamplingParams(
        temperature=0.7, 
        max_tokens=1024,
        top_p=0.9
    )

    results = []
    print(f"\n[Run] Starting Inference on {len(meta_data)} images...")
    
    for i in tqdm(range(0, len(meta_data), args.batch_size), desc="Processing Batches"):
        batch_meta = meta_data[i : i + args.batch_size]
        batch_inputs = prepare_vllm_inputs(batch_meta, processor)
        
        valid_inputs = [inp for inp in batch_inputs if inp is not None]
        valid_indices = [idx for idx, inp in enumerate(batch_inputs) if inp is not None]
        
        if not valid_inputs:
            continue
            
        outputs = llm.generate(valid_inputs, sampling_params=sampling_params, use_tqdm=False)
        
        for local_idx, out in enumerate(outputs):
            original_meta = batch_meta[valid_indices[local_idx]]
            generated_text = out.outputs[0].text

            parsed = parse_llm_output(generated_text)
            
            results.append({
                "filename": original_meta["filename"],
                "path": original_meta["path"],
                "label": parsed["label"],   # Suitable / Unsuitable
                "think": parsed["think"],  
                "raw_output": generated_text
            })

    total = len(results)
    unsuitable_count = sum(1 for r in results if r["label"] == "Unsuitable")
    suitable_count = sum(1 for r in results if r["label"] == "Suitable")
    error_count = total - unsuitable_count - suitable_count

    unsuitable_rate = (unsuitable_count / total * 100) if total > 0 else 0
    suitable_rate = (suitable_count / total * 100) if total > 0 else 0

    print("\n" + "="*60)
    print(f"AUDIT REPORT FOR: {input_path.name}")
    print("="*60)
    print(f"{'Total Images':<25}: {total}")
    print("-" * 60)
    print(f"{'UNSUITABLE (Violation)':<25}: {unsuitable_count}  ({unsuitable_rate:.2f}%)")
    print(f"{'SUITABLE (Safe)':<25}: {suitable_count}  ({suitable_rate:.2f}%)")
    print(f"{'Parse Errors':<25}: {error_count}")
    print("="*60)

    output_file = input_path / f"audit_result_{input_path.name}.json"
    try:
        with open(output_file, "w", encoding="utf-8") as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        print(f"\n[Done] Detailed JSON report saved to:\n-> {output_file}")
    except Exception as e:
        print(f"[Error] Could not save JSON: {e}")

if __name__ == "__main__":
    try:
        multiprocessing.set_start_method('spawn', force=True)
    except RuntimeError:
        pass
    
    main()