code_inference/rule_19_vllm.py

# -*- coding: utf-8 -*-
"""
使用方法:
python rule_19_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. Your job is to flag "Low-Quality / Amateur" advertising designs.
You have ZERO TOLERANCE for "Cheap Ad Styles" (often called "Niu Pi Xian" in Chinese context).

INPUT: One image and one natural-language question about design aesthetic and text harmony.

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").

OUTPUT FORMAT:
Return EXACTLY two blocks, no extra text:
<think>Detailed reasoning evaluating font effects, background integration, and aesthetic consistency against the 'cheap design' criteria...</think><answer>{"Answer": "<Suitable OR Unsuitable>", "Answer type": "Text-Design Harmony"}</answer>

=========================================
STRICT VIOLATION CRITERIA (If ANY match -> Unsuitable)
=========================================
1. **The "WordArt" Effect (廉价特效):**
   - **Bad Strokes:** Text uses heavy, amateurish strokes (thick white/colored outlines) that look jagged or pixelated.
   - **Fake 3D/Metal:** Outdated "Pseudo-3D" gradients (e.g., shiny gold/silver metal textures) that clash with a flat background.
   - **Cheap Glow:** Aggressive "Outer Glow" (neon glow) that makes the text look blurry or radioactive.
   - **Distortion:** Text is unprofessionally stretched, squeezed, or distorted strictly to fit a space.

2. **Visual Clutter & Conflict (背景冲突与拼贴感):**
   - **Legibility Loss:** Text is placed directly on top of a "Busy Photograph" (leaves, city streets, crowds) without a sufficient background mask, making it hard to read.
   - **Color Vibration:** Text color aggressively vibrates against the background (e.g., bright red text directly on bright green).
   - **Patchwork Style:** The text background looks like a "sticker" arbitrarily pasted onto a photo, completely ignoring the photo's lighting and perspective.

3. **Inconsistent Aesthetic (风格割裂):**
   - Foreground graphic elements (e.g., a cartoon/gaming style "Button" or "Banner") are superimposed on a realistic, high-res nature/human photograph. They do not belong in the same visual world.


=========================================
CRITERIA FOR 'SUITABLE' (NON-VIOLATION / GOOD DESIGN)
=========================================
1. **Clean Professionalism:** Professional typography with no cheap text effects (e.g., simple text like "xx折扣" is perfectly fine if the font is clean).
2. **Proper Integration:** Text placed on a solid, clean color background, or properly masked on a complex background.
3. **Cohesive Art Direction:** Clean, flat vector art that matches its surroundings visually.

=========================================
DECISION LOGIC
=========================================
- **Unsuitable**: If the design looks cheap, messy, outdated, features "WordArt" effects, or feels like a patched-together "Niu Pi Xian" ad.
- **Suitable**: If the design is clean, professional, and visually harmonious.
"""



SYS_PROMPT_TEXT="""You are a highly critical "Senior Art Director." Your goal is to evaluate the "Professional Polish" of splash ads. You have [ZERO TOLERANCE] for raw, unprocessed photos that look like amateur snapshots.

INPUT:
One image and one natural-language question about professional polish quality.

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").

OUTPUT FORMAT:
Return EXACTLY two blocks, no extra text:

<think>
Detailed reasoning evaluating lighting, color grading, and depth of field against the "passerby snapshot" criteria...
</think>
<answer>
{"Answer": "<Suitable OR Unsuitable>", "Answer type": "Professional Polish"}
</answer>

=========================================
STRICT VIOLATION CRITERIA (If ANY match -> Unsuitable)
=========================================
1. **Lack of Professional Post-Processing:**
   - The image appears to be a "Raw Photo" directly from a camera/phone without professional retouching.

2. **The "Amateur Snapshot" Aesthetic:**
   - The image looks like something a "passerby" could easily capture. It lacks the sophisticated framing, high-end texture, and artistic polish required for premium advertising.

3. **Absence of Value Conveyance:**
   - The image is visually "flat" and fails to evoke a sense of high quality. It does not use professional polish techniques to guide the viewer's emotions.

=========================================
CRITERIA FOR 'SUITABLE' (NON-VIOLATION / PREMIUM TEXTURE)
=========================================
1. **Professional Post-Processing:**
   - **Masterful Retouching:** The image shows clear evidence of professional professional polish (not straight-out-of-camera).

2. **The "High-End" Aesthetic:**
   - **Professionalism:** The image features a look that cannot be easily replicated by a passerby.
   - **Superior Texture:** Displays deliberate set design and artistic polish.
   
4. **Media Exemption:** Film stills or variety show photography are always classified as SAFE (Suitable).

=========================================
DECISION LOGIC
=========================================
- **Unsuitable**: If the image looks like an unprocessed, amateur snapshot with flat lighting and no professional polish polish.
- **Suitable**: If the image shows professional professional polish, or is a professional film/variety show still.
"""




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