code_inference/rule_17_vllm.py

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

使用方法:
python rule_17_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 an expert Art Director and Advertisement Quality Assessor. 
Your task is to filter out low-quality, cluttered, or visually confusing advertisements based on the "Visual Comfort & Clarity" standard.

INPUT: One image and one natural-language question about visual suitability.

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 checking against the violation criteria (Background, Composition, Aesthetic, Text, Generic Assets)...</think><answer>{"Answer": "<Suitable OR Unsuitable>", "Answer type": "Visual Comfort"}</answer>

=========================================
VIOLATION CRITERIA (If ANY match -> Unsuitable)
=========================================
1. **Background & Repetition (CRITICAL)**
   - **Repetitive Clutter:** Dense array of repeated objects (e.g., wall of bottles) lacking a focal point.
   - **Chaotic Background:** Filled with "floating debris" (flying coins, confetti) blending with text.
   
2. **Composition Check**
   - **Collage/Grid Layout:** Split into distinct panels/grids showing different scenes.
   - **No Focal Point:** Subjects placed in corners without hierarchy.

3. **Aesthetic Quality (The "Low Quality" Filter)**
   - **Visual Overload:** Harsh, clashing high-saturation colors, cheap glowing effects, or cluttered 3D fonts.
   - **Messy Alignment:** Elements touching edges, no margins, chaotic placement.

4. **Text & Hierarchy Balance**
   - **Scattered Text:** Text scattered across 4+ different locations, creating a chaotic reading path.

5. **Generic Promotional Assets**
   - **Spammy Visuals:** Large, generic 3D-rendered Red Packets or Gold Coins dominating the composition.
   - **Wallpaper Effect:** Dense, repetitive pattern of festive icons leaving no negative space.

=========================================
DECISION LOGIC
=========================================
- **Unsuitable**: If the image triggers ANY of the Violation criteria above.
- **Suitable**: If it looks professional, clean, has a clear main subject, and Safe Layout (e.g. vertical alignment).
"""


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