| |
| """ |
| |
| 使用方法: |
| 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"], |
| "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() |
