# -*- 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: Detailed reasoning checking against the violation criteria (Background, Composition, Aesthetic, Text, Generic Assets)...{"Answer": "", "Answer type": "Visual Comfort"} ========================================= 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'(.*?)', text, re.DOTALL) think_content = think_match.group(1).strip() if think_match else "" answer_match = re.search(r'(.*?)', 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()