Upload benchmark_load.py

benchmark_load.py

@@ -0,0 +1,354 @@
+import os
+import os.path as osp
+import pandas as pd
+import warnings
+from typing import Dict, List, Optional, Union
+import hashlib
+from tqdm import tqdm
+import urllib.request
+import json
+import base64
+from PIL import Image
+import io
+
+# 数据集配置
+DATASET_CONFIG = {
+    'MME': {
+        'url': 'https://opencompass.openxlab.space/utils/VLMEval/MME.tsv',
+        'md5': 'b36b43c3f09801f5d368627fb92187c3',
+        'type': 'Y/N'
+    },
+    'HallusionBench': {
+        'url': 'https://opencompass.openxlab.space/utils/VLMEval/HallusionBench.tsv',
+        'md5': '0c23ac0dc9ef46832d7a24504f2a0c7c',
+        'type': 'Y/N'
+    },
+    'POPE': {
+        'url': 'https://opencompass.openxlab.space/utils/VLMEval/POPE.tsv',
+        'md5': 'c12f5acb142f2ef1f85a26ba2fbe41d5',
+        'type': 'Y/N'
+    },
+    'AMBER': {
+        'url': 'https://huggingface.co/datasets/yifanzhang114/AMBER_base64/resolve/main/AMBER.tsv',
+        'md5': '970d94c0410916166e0a76ba75da7934',
+        'type': 'Y/N'
+    },
+    'MMBench_DEV_EN': {
+        'url': 'https://opencompass.openxlab.space/utils/benchmarks/MMBench/MMBench_DEV_EN.tsv',
+        'md5': 'b6caf1133a01c6bb705cf753bb527ed8',
+        'type': 'MCQ'
+    },
+    'SEEDBench2_Plus': {
+        'url': 'https://opencompass.openxlab.space/utils/benchmarks/SEEDBench/SEEDBench2_Plus.tsv',
+        'md5': 'e32d3216dc4f452b0fe497a52015d1fd',
+        'type': 'MCQ'
+    },
+    'ScienceQA_VAL': {
+        'url': 'https://opencompass.openxlab.space/utils/benchmarks/ScienceQA/ScienceQA_VAL.tsv',
+        'md5': '96320d05e142e585e7204e72affd29f3',
+        'type': 'MCQ'
+    },
+    'MMMU_TEST': {
+        'url': 'https://opencompass.openxlab.space/utils/VLMEval/MMMU_TEST.tsv',
+        'md5': 'c19875d11a2d348d07e5eb4bdf33166d',
+        'type': 'MCQ'
+    },
+    'OCRBench': {
+        'url': 'https://opencompass.openxlab.space/utils/VLMEval/OCRBench.tsv',
+        'md5': 'e953d98a987cc6e26ef717b61260b778',
+        'type': 'VQA'
+    },
+    'MathVista_MINI': {
+        'url': 'https://opencompass.openxlab.space/utils/VLMEval/MathVista_MINI.tsv',
+        'md5': 'f199b98e178e5a2a20e7048f5dcb0464',
+        'type': 'VQA'
+    },
+    'MathVision': {
+        'url': 'https://opencompass.openxlab.space/utils/VLMEval/MathVision.tsv',
+        'md5': '93f6de14f7916e598aa1b7165589831e',
+        'type': 'VQA'
+    },
+    'MMDU': {
+        'url': 'https://opencompass.openxlab.space/utils/VLMEval/MMDU.tsv',
+        'md5': '848b635a88a078f49aebcc6e39792061',
+        'type': 'MT'
+    },
+    'MIA-Bench': {
+        'url': 'https://opencompass.openxlab.space/utils/VLMEval/Mia-Bench.tsv',
+        'md5': '0b9de595f4dd40af18a69b94d89aba82',
+        'type': 'VQA'
+    }
+}
+
+class DownloadProgressBar(tqdm):
+    """下载进度条"""
+    def update_to(self, b=1, bsize=1, tsize=None):
+        if tsize is not None:
+            self.total = tsize
+        self.update(b * bsize - self.n)
+
+def download_file(url: str, filename: str) -> str:
+    """下载文件"""
+    try:
+        with DownloadProgressBar(unit='B', unit_scale=True, miniters=1, desc=filename) as t:
+            urllib.request.urlretrieve(url, filename=filename, reporthook=t.update_to)
+    except Exception as e:
+        warnings.warn(f'下载失败: {e}')
+        # 处理huggingface.co下载失败的情况
+        if 'huggingface.co' in url:
+            url_new = url.replace('huggingface.co', 'hf-mirror.com')
+            try:
+                download_file(url_new, filename)
+                return filename
+            except Exception as e:
+                warnings.warn(f'镜像下载也失败: {e}')
+                raise Exception(f'无法下载 {url}')
+        else:
+            raise Exception(f'无法下载 {url}')
+    
+    return filename
+
+def md5(file_path: str) -> str:
+    """计算文件MD5"""
+    hash_md5 = hashlib.md5()
+    with open(file_path, "rb") as f:
+        for chunk in iter(lambda: f.read(4096), b""):
+            hash_md5.update(chunk)
+    return hash_md5.hexdigest()
+
+def decode_base64_to_image_file(base64_str: str, output_path: str):
+    """将base64字符串解码为图片文件"""
+    # 移除data:image/jpeg;base64,前缀
+    if ',' in base64_str:
+        base64_str = base64_str.split(',')[1]
+    
+    # 解码base64
+    image_data = base64.b64decode(base64_str)
+    
+    # 保存图片
+    with open(output_path, 'wb') as f:
+        f.write(image_data)
+
+def read_ok(file_path: str) -> bool:
+    """检查文件是否可读"""
+    return os.path.exists(file_path) and os.path.getsize(file_path) > 0
+
+def toliststr(x):
+    """将输入转换为字符串列表"""
+    if isinstance(x, str):
+        return [x]
+    elif isinstance(x, list):
+        return [str(item) for item in x]
+    else:
+        return [str(x)]
+
+class DatasetProcessor:
+    """数据集处理器"""
+    
+    def __init__(self, data_root: str):
+        """初始化处理器"""
+        self.data_root = data_root
+        
+        os.makedirs(self.data_root, exist_ok=True)
+        self.img_root = os.path.join(self.data_root, 'images')
+        os.makedirs(self.img_root, exist_ok=True)
+    
+    def download_dataset(self, dataset_name: str) -> str:
+        """下载单个数据集"""
+        if dataset_name not in DATASET_CONFIG:
+            raise ValueError(f"不支持的数据集: {dataset_name}")
+        
+        config = DATASET_CONFIG[dataset_name]
+        url = config['url']
+        file_md5 = config['md5']
+        
+        # 构建文件路径
+        file_name = url.split('/')[-1]
+        data_path = os.path.join(self.data_root, file_name)
+        
+        # 检查文件是否已存在且MD5正确
+        if os.path.exists(data_path):
+            if md5(data_path) == file_md5:
+                print(f"✓ 数据集 {dataset_name} 已存在且MD5正确")
+                return data_path
+            else:
+                print(f"⚠ 数据集 {dataset_name} 存在但MD5不匹配，重新下载")
+        
+        # 下载文件
+        print(f"正在下载数据集 {dataset_name}...")
+        download_file(url, data_path)
+        
+        # 验证MD5
+        if md5(data_path) != file_md5:
+            raise ValueError(f"数据集 {dataset_name} MD5验证失败")
+        
+        print(f"✓ 数据集 {dataset_name} 下载成功")
+        return data_path
+    
+    def extract_images(self, dataset_name: str, data: pd.DataFrame) -> Dict[str, str]:
+        """提取数据集中的图像"""
+        dataset_img_root = os.path.join(self.img_root, dataset_name)
+        os.makedirs(dataset_img_root, exist_ok=True)
+        
+        image_paths = {}
+        
+        if 'image' in data.columns:
+            print(f"正在提取 {dataset_name} 的图像...")
+            for idx, row in tqdm(data.iterrows(), total=len(data), desc=f"提取 {dataset_name} 图像"):
+                index = row['index']
+                image_data = row['image']
+                
+                if pd.isna(image_data):
+                    continue
+                
+                # 处理图像数据
+                if isinstance(image_data, str) and len(image_data) > 64:
+                    # 假设是base64编码的图像
+                    image_path = os.path.join(dataset_img_root, f"{index}.jpg")
+                    if not read_ok(image_path):
+                        try:
+                            decode_base64_to_image_file(image_data, image_path)
+                        except Exception as e:
+                            print(f"⚠ 解码图像失败 (索引 {index}): {e}")
+                            continue
+                    image_paths[str(index)] = image_path
+                elif isinstance(image_data, list):
+                    # 处理多图像情况
+                    for i, img in enumerate(image_data):
+                        if isinstance(img, str) and len(img) > 64:
+                            image_path = os.path.join(dataset_img_root, f"{index}_{i+1}.jpg")
+                            if not read_ok(image_path):
+                                try:
+                                    decode_base64_to_image_file(img, image_path)
+                                except Exception as e:
+                                    print(f"⚠ 解码图像失败 (索引 {index}_{i+1}): {e}")
+                                    continue
+                            image_paths[f"{index}_{i+1}"] = image_path
+        
+        print(f"✓ 提取了 {len(image_paths)} 张图像")
+        return image_paths
+    
+    def process_dataset(self, dataset_name: str) -> Dict:
+        """处理单个数据集"""
+        print(f"\n=== 处理数据集: {dataset_name} ===")
+        
+        # 下载数据集
+        data_path = self.download_dataset(dataset_name)
+        
+        # 加载数据
+        data = pd.read_csv(data_path, sep='\t')
+        print(f"✓ 加载了 {len(data)} 个样本")
+        
+        # 提取图像
+        image_paths = self.extract_images(dataset_name, data)
+        
+        # 准备结果
+        config = DATASET_CONFIG[dataset_name]
+        result = {
+            'dataset_name': dataset_name,
+            'dataset_type': config['type'],
+            'total_samples': len(data),
+            'image_count': len(image_paths),
+            'data': data,
+            'image_paths': image_paths,
+            'columns': list(data.columns)
+        }
+        
+        # 添加样本数据
+        sample_data = []
+        for idx, row in data.head(3).iterrows():
+            sample = {
+                'index': row['index'],
+                'question': row.get('question', 'N/A'),
+                'answer': row.get('answer', 'N/A')
+            }
+            
+            # 添加选项（如果是MCQ类型）
+            if config['type'] == 'MCQ':
+                options = {}
+                for col in ['A', 'B', 'C', 'D', 'E']:
+                    if col in row and not pd.isna(row[col]):
+                        options[col] = row[col]
+                if options:
+                    sample['options'] = options
+            
+            # 添加图像路径
+            if str(row['index']) in image_paths:
+                sample['image_path'] = image_paths[str(row['index'])]
+            
+            sample_data.append(sample)
+        
+        result['sample_data'] = sample_data
+        
+        return result
+    
+    def process_datasets(self, dataset_names: List[str]) -> Dict[str, Dict]:
+        """处理多个数据集"""
+        if not dataset_names:
+            raise ValueError("数据集名称列表不能为空")
+        
+        # 验证数据集名称
+        invalid_datasets = [name for name in dataset_names if name not in DATASET_CONFIG]
+        if invalid_datasets:
+            raise ValueError(f"不支持的数据集: {invalid_datasets}")
+        
+        print(f"开始处理 {len(dataset_names)} 个数据集: {dataset_names}")
+        
+        results = {}
+        
+        for dataset_name in dataset_names:
+            try:
+                result = self.process_dataset(dataset_name)
+                results[dataset_name] = result
+                print(f"✓ 数据集 {dataset_name} 处理完成")
+            except Exception as e:
+                print(f"✗ 处理数据集 {dataset_name} 失败: {e}")
+                results[dataset_name] = None
+        
+        return results
+
+def process_vlmeval_datasets(dataset_names: List[str], data_root: str) -> Dict[str, Dict]:
+    """
+    主函数：处理VLMEval数据集
+    
+    Args:
+        dataset_names: 数据集名称列表，支持的数据集包括：
+            - Y/N类型: MME, HallusionBench, POPE, AMBER
+            - MCQ类型: MMBench_DEV_EN, SEEDBench2_Plus, ScienceQA_VAL, MMMU_TEST
+            - VQA类型: OCRBench, MathVista_MINI, MathVision, MIA-Bench
+            - MT类型: MMDU
+        data_root: 数据存储根目录
+    
+    Returns:
+        包含所有数据集处理结果的字典
+    """
+    processor = DatasetProcessor(data_root)
+    return processor.process_datasets(dataset_names)
+
+# 使用示例
+if __name__ == "__main__":
+    # 示例1：处理单个数据集
+    print("=== 示例1：处理单个数据集 ===")
+    result = process_vlmeval_datasets(['MMBench_DEV_EN'], data_root='/media/raid/workspace/jiangkailin/data_and_ckpt/dataset/cache/vlmeval')
+    
+    for dataset_name, dataset_result in result.items():
+        if dataset_result:
+            print(f"\n数据集: {dataset_name}")
+            print(f"类型: {dataset_result['dataset_type']}")
+            print(f"样本数: {dataset_result['total_samples']}")
+            print(f"图像数: {dataset_result['image_count']}")
+            print("前3个样本:")
+            for i, sample in enumerate(dataset_result['sample_data'], 1):
+                print(f"  样本{i}: 索引={sample['index']}, 问题={sample['question'][:50]}...")
+    
+    # # 示例2：处理多个数据集
+    # print("\n=== 示例2：处理多个数据集 ===")
+    # datasets = ['MME', 'POPE', 'OCRBench']
+    # results = process_vlmeval_datasets(datasets, data_root='/home/jiangkailin/mydisk/iclr26_evoke_dynamic_null_space/cache/vlmeval')
+    
+    # print(f"\n处理完成，共处理 {len(results)} 个数据集:")
+    # for dataset_name, dataset_result in results.items():
+    #     if dataset_result:
+    #         print(f"  ✓ {dataset_name}: {dataset_result['total_samples']} 样本, {dataset_result['image_count']} 图像")
+    #     else:
+    #         print(f"  ✗ {dataset_name}: 处理失败")