| |
| """ |
| 使用方法: |
| python rule_14_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. Crucial Context: You must distinguish between "Aggressive E-commerce Marketing" (Professional) and "Amateur Sloppiness" (Violation). High-resolution assets, standard platform badges, and professional 3D renders are SAFE. |
| |
| INPUT: One image and one natural-language question about design aesthetic and text harmony. |
| |
| YOUR TASK: |
| |
| Determine if the image is a VIOLATION (Unsuitable) or SAFE (Suitable) based on the criteria below. |
| |
| 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...</think><answer>{"Answer": "<Suitable OR Unsuitable>", "Answer type": "Text Design Harmony"}</answer> |
| |
| ========================================= STRICT VIOLATION CRITERIA (If ANY match -> Unsuitable) |
| The "WordArt" Effect (廉价特效): |
| |
| Technical Failure: ONLY flag if text is pixelated, jagged, or uses 1990s-style rainbow/neon gradients. |
| |
| Distortion: Text is unprofessionally stretched or squeezed (breaking the font's aspect ratio). |
| |
| Amateur Strokes: Thick, vibrating outlines that look like they were made in MS Paint, not professional design software. |
| |
| Note: High-res 3D fonts, clean gold textures, and smooth gradients are PROFESSIONAL and SAFE. |
| |
| Visual Clutter & Conflict (背景冲突与拼贴感): |
| |
| Resolution Mismatch: A low-res/blurry graphic pasted onto a high-res photo. |
| |
| Zero Integration: Elements that have NO shadows, NO lighting consistency, and look like accidental "floating" errors. |
| |
| Legibility Loss: Text is truly unreadable due to background chaos without any masking. |
| |
| Note: Standard UI elements (Pill buttons, Price tags, Promo badges like "百亿补贴") are INTENTIONAL overlays and are SAFE. |
| |
| Inconsistent Aesthetic (风格割裂): |
| |
| Flag ONLY if elements are accidentally mismatched (e.g., a hand-drawn sketch randomly appearing in a high-tech 3D render without stylistic intent). |
| |
| Note: 3D mascots or cartoon characters placed on realistic backgrounds for marketing purposes are a VALID style and are SAFE. |
| |
| ========================================= CRITERIA FOR 'SUITABLE' (NON-VIOLATION / GOOD DESIGN) |
| Commercial Execution: High-resolution assets, clean font edges, and professional lighting/shadows. |
| |
| Platform Legitimacy: Presence of brand logos (Alipay, Taobao, Banks, China Gold) and standard e-commerce UI components. |
| |
| Intentional Hierarchy: Even if the design is "loud" (Red/Gold), it is Suitable if the text is aligned and the layout is purposeful. |
| |
| ========================================= DECISION LOGIC |
| Unsuitable: If the design shows Technical Failure (pixelation, distortion, 90s-style WordArt) or looks like a non-designer's mistake. |
| |
| Suitable: If the design follows Commercial Logic (Standard e-commerce banners, High-res renders, Professional marketing layouts). When in doubt, if the image looks like it's from a major App, it is SUITABLE. |
| """ |
|
|
|
|
|
|
| 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"], |
| "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() |
