opencompass/TextEdit

1#!/usr/bin/env python32import os3import sys4import time5 6# ==================== 1. Strict environment variable lock (avoid CPU oversubscription deadlocks) ====================7# Must be set before importing torch/numpy8os.environ["OMP_NUM_THREADS"] = "1"9os.environ["MKL_NUM_THREADS"] = "1"10os.environ["OPENBLAS_NUM_THREADS"] = "1"11os.environ["VECLIB_MAXIMUM_THREADS"] = "1"12os.environ["NUMEXPR_NUM_THREADS"] = "1"13os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"14# ==============================================================================15 16from torch.distributed.elastic.multiprocessing.errors import record17import argparse18import json19import re20from pathlib import Path21from typing import Dict, List, Tuple, Optional22import logging23import numpy as np24from PIL import Image25from tqdm import tqdm26import torch27import collections28import shutil29import tempfile30import copy31import unicodedata32import torch.multiprocessing as mp33from classic_ocr_tools import norm, poly_to_bbox, bbox_iou, point_in_bbox, bbox_center, union_bboxes, group_into_lines34 35# ==================== tools function ====================36 37def convert_numpy_types(obj):38    """Recursively convert numpy types to Python native types for JSON serialization"""39    if isinstance(obj, np.integer):40        return int(obj)41    elif isinstance(obj, np.floating):42        return float(obj)43    elif isinstance(obj, np.ndarray):44        return obj.tolist()45    elif isinstance(obj, dict):46        return {key: convert_numpy_types(value) for key, value in obj.items()}47    elif isinstance(obj, list):48        return [convert_numpy_types(item) for item in obj]49    return obj50 51 52class UnifiedMetricsEvaluator:53    def __init__(self, device: str = "auto", cache_dir: str = None):54        """Initialize evaluator"""55        self.device = "cuda" if torch.cuda.is_available() and device == "auto" else device56        self.models = {}57        self.cache_dir = cache_dir58        self._load_models()59    60    def _load_models(self):61        """Load all required models (Lazy Import Mode)"""62        print(f"[{self.device}] Loading models...")63        64        # 1. PaddleOCR65        try:66            from paddleocr import PaddleOCR67            # use_gpu=False to avoid multi-process GPU memory contention and hangs68            self.models['ocr'] = PaddleOCR(use_angle_cls=True, lang="en", show_log=False, use_gpu=False)69            self.paddleocr_available = True70            print(f"[{self.device}] PaddleOCR initialized")71        except ImportError:72            self.paddleocr_available = False73            logging.warning("PaddleOCR not available")74        75        # 2. CLIP76        try:77            import clip78            clip_model, clip_preprocess = clip.load("ViT-L/14", device=self.device, jit=False, download_root=self.cache_dir)79            clip_model.eval()80            self.models['clip_official'] = clip_model81            self.models['clip_official_preprocess'] = clip_preprocess82            self.clip_available = True83            print(f"[{self.device}] CLIP loaded")84        except Exception as e:85            logging.error(f"Failed to load CLIP: {e}")86            self.clip_available = False87        88        # 3. OpenCLIP & Aesthetic89        try:90            import open_clip91            model, _, preprocess = open_clip.create_model_and_transforms('ViT-L-14', pretrained="openai", cache_dir=self.cache_dir)92            model.to(self.device)93            model.eval()94            self.models['openclip'] = model95            self.models['openclip_preprocess'] = preprocess96            97            aesthetic_model = self._load_aesthetic_model()98            if aesthetic_model:99                self.models['aesthetic'] = aesthetic_model100                print(f"[{self.device}] Aesthetic model loaded")101            self.openclip_available = True102        except Exception as e:103            self.openclip_available = False104 105    def _load_aesthetic_model(self):106        try:107            import torch.nn as nn108            model_path = os.path.join(self.cache_dir, "sa_0_4_vit_l_14_linear.pth")109            if os.path.exists(model_path):110                m = nn.Linear(768, 1)111                s = torch.load(model_path, map_location=self.device)112                m.load_state_dict(s)113                m.eval().to(self.device)114                return m115        except Exception as e:116            logging.warning(f"Could not load aesthetic model: {e}")117        return None118 119    def get_ld(self, ls1: str, ls2: str) -> float:120        """Calculate normalized version of Levenshtein distance"""121        if not self.paddleocr_available:122            return 0.0123        import Levenshtein124        edit_dist = Levenshtein.distance(ls1, ls2)125        return 1 - edit_dist / (max(len(ls1), len(ls2)) + 1e-5)126 127    def compute_roi_ned(self, raw_items, gen_items, source_text, target_text) -> float:128        """Compute ROI-based normalized edit distance (ROI-aware NED)."""129        import Levenshtein130        source_norm = norm(source_text)131 132        best_sim = 0.0133        source_bbox = None134 135        for item in raw_items:136            sim = self.get_ld(source_norm, norm(item["text"]))137            if sim > 0.6 and sim > best_sim:138                best_sim = sim139                source_bbox = item["bbox"]140 141        if source_bbox is None:142            return 0.0143 144        pred_texts_in_roi = [145            item for item in gen_items146            if bbox_iou(item["bbox"], source_bbox) > 0.3147        ]148        if not pred_texts_in_roi:149            return 0.0150 151        pred_texts_in_roi.sort(key=lambda x: x["bbox"][0])152 153        pred_norm = norm("".join([it["text"] for it in pred_texts_in_roi]))154        target_norm = norm(target_text)155 156        ned_score = self.get_ld(target_norm, pred_norm)157 158        if self.get_ld(source_norm, pred_norm) > 0.9 and self.get_ld(source_norm, target_norm) < 0.5:159            ned_score *= 0.2160 161        return ned_score162 163    def compute_ocr_metrics_textedit(164        self,165        raw_img_path: str,166        gen_img_path: str,167        source_text: str,168        target_text: str,169        target_weight: float = 0.5,170        iou_threshold: float = 0.5171    ) -> Dict:172        """OCR evaluation for text editing."""173        default_res = {174            "target_accuracy": 0.0,175            "precision": 0.0,176            "recall": 0.0,177            "f1": 0.0,178            "roi_ned": 0.0,179        }180 181        if not self.paddleocr_available or "ocr" not in self.models:182            return default_res183 184        try:185            source_norm = norm(source_text)186            target_norm = norm(target_text)187 188            raw_img = Image.open(raw_img_path)189            gen_img_original = Image.open(gen_img_path)190 191            # Uniformly resize generated image to original size for fair comparison192            if gen_img_original.mode != "RGB":193                gen_img_original = gen_img_original.convert("RGB")194 195            if gen_img_original.size != raw_img.size:196                gen_img_resized = gen_img_original.resize(raw_img.size, Image.LANCZOS)197            else:198                gen_img_resized = gen_img_original199 200            # OCR: raw201            raw_ocr = self.models["ocr"].ocr(raw_img_path, cls=True)202            raw_lines_src = raw_ocr[0] if (raw_ocr and raw_ocr[0]) else []203            raw_items = [204                {"text": line[1][0], "bbox": poly_to_bbox(line[0]),205                 'poly': line[0], 'score': line[1][1]}206                for line in raw_lines_src207            ]208            if not raw_items:209                return default_res210 211            raw_lines = group_into_lines(raw_items)212 213            best_line_idx = -1214            best_sim = 0.0215            for i, line_items in enumerate(raw_lines):216                line_text = "".join(item["text"] for item in line_items)217                line_norm = norm(line_text)218                219                if source_norm in line_norm:220                    sim = 1.0221                else:222                    sim = self.get_ld(source_norm, line_norm)223                224                if sim > best_sim:225                    best_sim = sim226                    best_line_idx = i227 228            if best_line_idx == -1 or best_sim < 0.5:229                return default_res230 231            # OCR: gen (numpy array)232            gen_ocr = self.models["ocr"].ocr(np.array(gen_img_resized), cls=True)233            gen_lines_src = gen_ocr[0] if (gen_ocr and gen_ocr[0]) else []234            gen_items = [235                {"text": line[1][0], "bbox": poly_to_bbox(line[0]),236                 'poly': line[0], 'score': line[1][1]}237                for line in gen_lines_src238            ]239            if not gen_items:240                return default_res241 242            # Region split243            raw_edit_line = raw_lines[best_line_idx]244            raw_edit_region = union_bboxes([item['bbox'] for item in raw_edit_line])245            gen_edit_region = raw_edit_region246            247            gen_region_items = []248            gen_bg_items = []249            250            for item in gen_items:251                iou = bbox_iou(item['bbox'], gen_edit_region)252                center = bbox_center(item['bbox'])253                if iou > iou_threshold or point_in_bbox(center, gen_edit_region):254                    gen_region_items.append(item)255                else:256                    gen_bg_items.append(item)257 258            # Target similarity: find target text in the edited region259            best_target_sim = 0.0260            if gen_region_items:261                for it in gen_region_items:262                    best_target_sim = max(best_target_sim, self.get_ld(target_norm, norm(it["text"])))263 264                for line in group_into_lines(gen_region_items):265                    text_merged = "".join(it["text"] for it in line)266                    best_target_sim = max(best_target_sim, self.get_ld(target_norm, norm(text_merged)))267 268            # Penalize if source still exists while target does not269            gen_all_text_norm = norm("".join(item["text"] for item in gen_items))270            if source_norm in gen_all_text_norm and target_norm not in gen_all_text_norm:271                best_target_sim *= 0.2272 273            target_accuracy = best_target_sim274 275            # Background line alignment similarity276            raw_bg_items = [277                item278                for i, line in enumerate(raw_lines)279                if i != best_line_idx280                for item in line281            ]282            raw_bg_lines = group_into_lines(raw_bg_items)283            gen_bg_lines = group_into_lines(gen_bg_items)284 285            used_gen_indices = set()286            bg_sims = []287            288            for raw_line in raw_bg_lines:289                raw_l_norm = norm(''.join(it['text'] for it in raw_line))290                current_best_sim = 0.0291                best_gen_idx = -1292                293                for j, gen_line in enumerate(gen_bg_lines):294                    if j in used_gen_indices:295                        continue296                    gen_l_norm = norm(''.join(it['text'] for it in gen_line))297                    sim = self.get_ld(raw_l_norm, gen_l_norm)298                    if sim > current_best_sim:299                        current_best_sim = sim300                        best_gen_idx = j301                302                if best_gen_idx != -1:303                    used_gen_indices.add(best_gen_idx)304                305                bg_sims.append(current_best_sim)306 307            # Compute Precision / Recall / F1308            w_target = target_weight309            num_gt_bg = len(raw_bg_lines)310            num_pred_bg = len(gen_bg_lines)311            w_bg_unit = (1.0 - w_target) / num_gt_bg if num_gt_bg > 0 else 0.0312            313            tp_score = (target_accuracy * w_target) + (sum(bg_sims) * w_bg_unit)314            gt_total_weight = w_target + (num_gt_bg * w_bg_unit)315            pred_total_weight = w_target + (num_pred_bg * w_bg_unit)316            317            recall = tp_score / gt_total_weight if gt_total_weight > 0 else 0.0318            precision = tp_score / pred_total_weight if pred_total_weight > 0 else 0.0319            f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0320 321            roi_ned_score = self.compute_roi_ned(322                raw_items,323                gen_items,324                source_text,325                target_text,326            )327 328            return {329                "target_accuracy": target_accuracy,330                'precision': precision,331                'recall': recall,332                "f1": f1,333                "roi_ned": roi_ned_score,334            }335 336        except Exception as e:337            logging.error(f"Compute OCR metrics failed: {e}")338            return default_res339 340    def compute_clip_score_batch(self, image_paths: List[str], texts: List[str]) -> List[float]:341        """Batch compute CLIPScore"""342        if not self.clip_available or 'clip_official' not in self.models:343            return [0.0] * len(image_paths)344        345        try:346            import clip347            from sklearn.preprocessing import normalize348            import sklearn.preprocessing349            from packaging import version350            import warnings351            352            processed_texts = []353            for text in texts:354                processed_texts.append(text)355            356            images = []357            for image_path in image_paths:358                try:359                    image = Image.open(image_path)360                    image_input = self.models['clip_official_preprocess'](image).unsqueeze(0)361                    images.append(image_input)362                except Exception as e:363                    logging.warning(f"Failed to load image {image_path}: {e}")364                    images.append(torch.zeros(1, 3, 224, 224))365            366            if not images:367                return []368 369            images_batch = torch.cat(images, dim=0).to(self.device)370            texts_batch = clip.tokenize(processed_texts, truncate=True).to(self.device)371            372            with torch.no_grad():373                image_features = self.models['clip_official'].encode_image(images_batch)374                text_features = self.models['clip_official'].encode_text(texts_batch)375                376                image_features_np = image_features.cpu().numpy()377                text_features_np = text_features.cpu().numpy()378                379                if version.parse(np.__version__) < version.parse('1.21'):380                    image_features_np = sklearn.preprocessing.normalize(image_features_np, axis=1)381                    text_features_np = sklearn.preprocessing.normalize(text_features_np, axis=1)382                else:383                    warnings.warn('Using compat normalization')384                    image_features_np = image_features_np / np.sqrt(np.sum(image_features_np**2, axis=1, keepdims=True))385                    text_features_np = text_features_np / np.sqrt(np.sum(text_features_np**2, axis=1, keepdims=True))386                387                similarities = np.sum(image_features_np * text_features_np, axis=1)388                clip_scores = 2.5 * np.clip(similarities, 0, None)389                390                return clip_scores.tolist()391        392        except Exception as e:393            logging.error(f"Batch CLIP score computation failed: {e}")394            return [0.0] * len(image_paths)395 396    def compute_aesthetic_score(self, image_path: str) -> float:397        """Calculate aesthetic score"""398        if not self.openclip_available or 'aesthetic' not in self.models or 'openclip' not in self.models:399            return 0.0400        401        try:402            image = Image.open(image_path)403            image_input = self.models['openclip_preprocess'](image).unsqueeze(0).to(self.device)404            405            with torch.no_grad():406                image_features = self.models['openclip'].encode_image(image_input)407                image_features /= image_features.norm(dim=-1, keepdim=True)408                prediction = self.models['aesthetic'](image_features)409                410            return prediction.cpu().numpy().item()411        412        except Exception as e:413            logging.error(f"Aesthetic score computation failed for {image_path}: {e}")414            return 0.0415 416    def evaluate_group(self, model_name: str, group_name: str, jsonl_files: List[str], 417                       benchmark_dir: str, gt_root_dir: str, model_output_root: str) -> Dict:418        """Evaluate a group of files for a specific model"""419 420        # Aggregated Metrics421        agg_metrics = {422            'target_accuracies': [],423            'precisions': [],424            'recalls': [],425            'f1s': [],426            'roi_neds': [],427            'clip_scores': [],428            'aesthetic_scores': [],429            'image_count': 0430        }431 432        # Store per-sample detailed results433        detailed_results = []434        print(f"[{self.device}] Processing {model_name} | Group: {group_name} ({len(jsonl_files)} files)")435 436        for jsonl_file in jsonl_files:437            full_path = os.path.join(benchmark_dir, jsonl_file)438            if not os.path.exists(full_path):439                logging.warning(f"File not found: {full_path}")440                continue441 442            with open(full_path, 'r', encoding='utf-8') as f:443                data_entries = [json.loads(line) for line in f if line.strip()]444            445            # 1. Prepare Batches for this file446            batch_paths, batch_prompts, valid_indices = [], [], []447 448            for idx, entry in enumerate(data_entries):449                # Resolve Model Generated Image Path450                raw_img_rel = entry.get('original_image', '')451                # class_id = parts[0]   # e.g., '1.1.1'452                # filename = parts[-1]  # e.g., '1901685000029.0.jpg'453                gen_img_path = os.path.join(model_output_root, model_name, raw_img_rel.split('/')[0], raw_img_rel.split('/')[-1])454 455                if gen_img_path and os.path.exists(gen_img_path):456                    # For CLIP/Aesthetic/OCR457                    batch_paths.append(gen_img_path)458                    batch_prompts.append(entry.get('gt_caption', ''))459                    valid_indices.append(idx)460            461            # 2. Compute CLIP Batch462            if batch_paths:463                clip_scores = self.compute_clip_score_batch(batch_paths, batch_prompts)464                clip_map = {idx: score for idx, score in zip(valid_indices, clip_scores)}465            else:466                clip_map = {}467 468            # 3. Compute Per-Image Metrics469            for idx, entry in enumerate(tqdm(data_entries, desc=f"  Evaluating {jsonl_file}", leave=False)):470                if idx not in clip_map:471                    continue472                473                raw_img_rel  = entry.get('original_image', '')474                raw_img_path = os.path.join(gt_root_dir, raw_img_rel)475                gt_img_rel   = entry.get('gt_image', '')476                gt_img_path  = os.path.join(gt_root_dir, gt_img_rel)477                gen_img_path = os.path.join(model_output_root, model_name, raw_img_rel.split('/')[0], raw_img_rel.split('/')[-1])478 479                # Read source_text and target_text directly from entry480                source_text = entry.get('source_text', '')481                target_text = entry.get('target_text', '')482 483                # CLIPScore484                clip_score = clip_map[idx]485                agg_metrics['clip_scores'].append(clip_score)486                487                # Aesthetic Score488                aesthetic_score = self.compute_aesthetic_score(gen_img_path)489                agg_metrics['aesthetic_scores'].append(aesthetic_score)490 491                # OCR metrics (pass source_text and target_text)492                ocr_res = self.compute_ocr_metrics_textedit(493                    raw_img_path=raw_img_path,494                    gen_img_path=gen_img_path,495                    source_text=source_text,496                    target_text=target_text,497                    target_weight=0.5,498                    iou_threshold=0.5499                )500 501                # Aggregate metrics502                agg_metrics['image_count'] += 1503                agg_metrics['target_accuracies'].append(ocr_res['target_accuracy'])504                agg_metrics['precisions'].append(ocr_res['precision'])505                agg_metrics['recalls'].append(ocr_res['recall'])506                agg_metrics['f1s'].append(ocr_res['f1'])507                agg_metrics['roi_neds'].append(ocr_res['roi_ned'])508                509                # Save per-sample detailed information510                detailed_results.append({511                    'id': entry.get('id'),512                    'prompt': entry.get('prompt', ''),513                    'path': {514                        'original_image': raw_img_path,515                        'edited_image': gen_img_path,516                        'gt_image': gt_img_path517                    },518                    'score': {519                        'ocr_accuracy': float(ocr_res['target_accuracy']),520                        'ocr_precision': float(ocr_res['precision']),521                        'ocr_recall': float(ocr_res['recall']),522                        'ocr_f1': float(ocr_res['f1']),523                        'clip_score': float(clip_score),524                        'ned_score': float(ocr_res['roi_ned']),525                        'aesthetic_score': float(aesthetic_score)526                    }527                })528 529        def sm(l): 530            return np.mean(l) if l else 0.0531        final_group_results = {532            'Group': group_name,533            'OCR Accuracy': sm(agg_metrics['target_accuracies']),534            'OCR Precision': sm(agg_metrics['precisions']),535            'OCR Recall': sm(agg_metrics['recalls']),536            'OCR F1': sm(agg_metrics['f1s']),537            'CLIPScore': sm(agg_metrics['clip_scores']),538            'NED': sm(agg_metrics['roi_neds']),539            'Aesthetic Score': sm(agg_metrics['aesthetic_scores']),540            'Total Images': agg_metrics['image_count'],541            'detailed_results': detailed_results542        }543        544        return final_group_results545 546def worker_process(model_name, gpu_id, args):547    try:548        device = f"cuda:{gpu_id}"549        print(f"\n>>> Worker started: Model={model_name} on Device={device}")550        os.environ["TORCH_HOME"] = args.cache_dir 551        evaluator = UnifiedMetricsEvaluator(device=device, cache_dir=args.cache_dir)552        553        groups = {554            "Virtual": [555                "1.1.1.jsonl", "1.1.2.jsonl", "1.1.3.jsonl",556                "1.2.1.jsonl", "1.2.2.jsonl",557                "1.3.1.jsonl", "1.3.2.jsonl",558                "1.4.1.jsonl", "1.4.2.jsonl", "1.4.3.jsonl", "1.4.4.jsonl"559            ],560            "Real": [561                "2.1.jsonl", "2.2.jsonl", "2.3.jsonl", "2.4.jsonl", 562                "2.5.jsonl", "2.6.jsonl", "2.7.jsonl"563            ]564        }565 566        # Evaluate567        res_v = evaluator.evaluate_group(model_name, "Virtual", groups["Virtual"], args.benchmark_dir, args.gt_root_dir, args.model_output_root)568        res_r = evaluator.evaluate_group(model_name, "Real", groups["Real"], args.benchmark_dir, args.gt_root_dir, args.model_output_root)569        570        metrics_order = [571                "OCR Accuracy", "OCR Precision", "OCR Recall", "OCR F1",572                "NED", "CLIPScore", "Aesthetic Score"573            ]574        final_res = {575            'summary_by_model': {model_name: {'Virtual': res_v, 'Real': res_r}}, 576            'metrics_list': metrics_order577            }578        579        # Print result table580        print(f"\n--- Results for {model_name} ---")581        print(f"{'Metric':<20} | {'Real':<15} | {'Virtual':<15}")582        print("-" * 60)583        for m in final_res['metrics_list']:584            print(f"{m:<20} | {res_r.get(m, 0.0):.4f}          | {res_v.get(m, 0.0):.4f}")585        print("-" * 60)586 587        out_path = os.path.join(args.output_dir, f"{model_name}.json")588        with open(out_path, 'w') as f:589            json.dump(convert_numpy_types(final_res), f, indent=4, ensure_ascii=False)590 591        print(f">>> Worker finished: {model_name}. Saved to {out_path}")592    except Exception as e:593        print(f"!!! Error in worker: {e}"); import traceback; traceback.print_exc()594 595@record596def main():597    parser = argparse.ArgumentParser(description='Unified text-to-image generation evaluation tool')598    parser.add_argument('--benchmark_dir', required=True, help='Directory containing the .jsonl files')599    parser.add_argument('--gt_root_dir', required=True, help='Root directory for Ground Truth images')600    parser.add_argument('--model_output_root', required=True, help='Root directory where model outputs are stored')601    parser.add_argument('--output_dir', required=True, help='result output file path')602    parser.add_argument('--models', default='bagel', help='Comma separated list of model names')603    parser.add_argument('--cache_dir', required=True, help='HuggingFace model cache directory path')604    args = parser.parse_args()605    606    if args.cache_dir:607        os.environ['TORCH_HOME'] = args.cache_dir608    mp.set_start_method('spawn', force=True)609    model_list = [m.strip() for m in args.models.split(',') if m.strip()]610    gpu_count = torch.cuda.device_count() or 1611    612    procs = []613    for i, model in enumerate(model_list):614        p = mp.Process(target=worker_process, args=(model, i % gpu_count, args))615        p.start()616        procs.append(p)617        time.sleep(30)618 619    for p in procs:620        p.join()621 622if __name__ == "__main__":623    main()