code_inference/rule_18_vllm.py

# -*- coding: utf-8 -*-
"""

使用方法:
python rule_18_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:
Analyze the image to determine if the layout violates professional "Composition & Spacing" standards. 

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": "Layout Breathability Check"}</answer>

=========================================
STRICT VIOLATION CRITERIA (If ANY match -> Unsuitable)
=========================================
1. **Lack of Breathing Room (Core Crowding Violation):**
   - Core Violation: The main subject (the specific product or hero character, excluding the background image), the headline text, and the logo are placed too close to each other.
   - Exclusions: This criterion does not apply to text appearing physically on the product packaging or artistic fonts that are visually integrated into the product itself.
   - The "Small Print" Nuance: While major design modules require significant negative space, secondary small text (such as annotations or footnotes) is permitted to have smaller gaps relative to other elements. However, these small characters must not be too close to other elements or edges; they must maintain a basic visual distance to avoid a sense of "clinging," "tangency," or extreme squeezing.
   - Visual Feel: The overall design feels "heavy" or "claustrophobic" because major modules lack sufficient negative space between them.

2. **Edge Tension (贴边风险):**
   - Elements are "touching" or "tangent" to each other or the border without intentional overlapping.

3. **Information Overload (信息堆砌):**
   - The layout is filled with too many text blocks or icons with no clear separation. 
   - There is no clear "Visual Path"; the eye doesn't know where to rest because every element is competing for attention and space simultaneously.

=========================================
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()