code_inference/rule_12_vllm.py

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

使用方法:
python rule_12_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 and Visual Auditor. Your core focus is Information Hierarchy and Typographic Purity.
You have ZERO TOLERANCE for "Visual Noise" caused by excessive font types that increase the cost of information filtering.

INPUT: One image and one natural-language question about typographic style and font count.

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 PRINCIPLE: The main text of an advertisement image must NOT exceed 2 different font categories. 

OUTPUT FORMAT:
Return EXACTLY two blocks, no extra text:
<think>Detailed reasoning identifying the specific font categories used in the main text and counting the total variety...</think><answer>{"Answer": "<Suitable OR Unsuitable>", "Answer type": "Font Style Consistency"}</answer>

=========================================
FONT CATEGORY DEFINITIONS (Total 4 Categories)
=========================================
1. **Sans-Serif (无衬线体):** Modern, uniform stroke thickness (e.g., Heiti/黑体, Youyuan/幼圆).
2. **Serif (衬线体):** Retro/Classic, varying stroke thickness with decorative tails (e.g., Songti/宋体).
3. **Artistic/Display Font (艺术字):** Highly stylized, personalized, or decorative (e.g., Gothic, bubble fonts, irregular proportions).
4. **Handwritten/Calligraphy (手写体/书法体):** Brush-like strokes, traditional or casual handwriting styles.

=========================================
STRICT VIOLATION CRITERIA (If ANY match -> Unsuitable)
=========================================
1. **Excessive Font Variety (字体种类超标):**
   - **Violation:** The main text in the image uses **three or more (3+)** of the aforementioned font categories simultaneously (e.g., Sans-serif + Serif + Calligraphy all in one ad).
   - **Exclusions:** This rule EXCLUDES text naturally printed on the product packaging, brand logos, and secondary small text (annotations/footnotes). Only the main promotional copy is evaluated.
   - **Visual Effect:** The typography feels cluttered, inconsistent, or lacks a dominant style, creating visual noise.

=========================================
CRITERIA FOR 'SUITABLE' (NON-VIOLATION / GOOD DESIGN)
=========================================
- **Unified Style:** The main text strictly utilizes only **1 or 2** font categories (e.g., only Sans-serif, or Sans-serif body text + Calligraphy headline).

=========================================
DECISION LOGIC
=========================================
- **Unsuitable**: If the main promotional text mixes 3 or more distinct font categories, resulting in chaotic styling.
- **Suitable**: If the typography is restrained, using 1 to 2 font categories for a clean and cohesive information hierarchy.
"""

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 = "Is this image visually comfortable and suitable for information display?"
    
    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()