insomnia7/gpcv_incontext_bench

1import os2import json3from PIL import Image4from pathlib import Path5 6def yolo_to_xyxy(bbox, img_width, img_height):7    """8    将 YOLO 格式 (center_x, center_y, width, height) 转换为 xyxy 格式9    坐标是归一化的，需要转换为绝对坐标10    """11    center_x, center_y, width, height = bbox12    x1 = (center_x - width / 2) * img_width13    y1 = (center_y - height / 2) * img_height14    x2 = (center_x + width / 2) * img_width15    y2 = (center_y + height / 2) * img_height16    return [int(x1), int(y1), int(x2), int(y2)]17 18def process_folder(category_path, category_name, output_file):19    """20    处理单个类别文件夹21    """22    images_dir = category_path / "images"23    labels_dir = category_path / "labels_yolo"24    25    if not images_dir.exists() or not labels_dir.exists():26        print(f"警告: {category_path} 中缺少 images 或 labels_yolo 文件夹")27        return 0, 028    29    # 获取所有图片文件（支持更多格式）30    image_files = []31    for ext in ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.tif', '*.tiff', '*.JPG', '*.PNG', '*.TIF', '*.TIFF']:32        image_files.extend(images_dir.glob(ext))33    34    processed_count = 035    total_bboxes = 036    37    for img_path in image_files:38        # 获取对应的标签文件39        label_file = labels_dir / f"{img_path.stem}.txt"40        if not label_file.exists():41            continue42        43        # 获取图片尺寸44        try:45            with Image.open(img_path) as img:46                img_width, img_height = img.size47        except Exception as e:48            print(f"错误: 无法读取图片 {img_path}, {e}")49            continue50        51        # 读取标签文件52        bboxes = []53        with open(label_file, 'r') as f:54            for line in f:55                line = line.strip()56                if not line:57                    continue58                parts = line.split()59                if len(parts) != 5:60                    continue61                62                class_id = int(parts[0])63                center_x = float(parts[1])64                center_y = float(parts[2])65                width = float(parts[3])66                height = float(parts[4])67                68                # 转换为 xyxy 格式69                xyxy = yolo_to_xyxy([center_x, center_y, width, height], img_width, img_height)70                71                bboxes.append({72                    "bbox": xyxy,  # [x1, y1, x2, y2]73                    "category": category_name,74                    "class_id": class_id75                })76        77        # 构造 JSONL 行78        json_line = {79            "image": str(img_path.absolute()),80            "image_path": str(img_path),81            "width": img_width,82            "height": img_height,83            "bboxes": bboxes84        }85        86        # 写入输出文件87        output_file.write(json.dumps(json_line, ensure_ascii=False) + '\n')88        processed_count += 189        total_bboxes += len(bboxes)90    91    return processed_count, total_bboxes92 93def main():94    # 当前目录95    base_dir = Path("/home/disk2/hjl/ICL_QWEN/ICL_benchmark")96    97    # 需要处理的文件夹（排除 dinov3 开头的）98    exclude_prefixes = ["dinov3"]99    100    # 输出文件101    output_path = base_dir / "dataset.jsonl"102    103    total_images = 0104    total_bboxes_all = 0105    106    with open(output_path, 'w', encoding='utf-8') as outfile:107        # 遍历所有文件夹108        for item in base_dir.iterdir():109            if not item.is_dir():110                continue111            112            # 检查是否需要排除113            should_exclude = False114            for prefix in exclude_prefixes:115                if item.name.startswith(prefix):116                    should_exclude = True117                    break118            119            if should_exclude:120                print(f"跳过文件夹: {item.name}")121                continue122            123            # 处理该类别文件夹124            print(f"处理类别: {item.name}", end=' ', flush=True)125            processed, bboxes = process_folder(item, item.name, outfile)126            print(f"完成 - 处理了 {processed} 张图片, {bboxes} 个标注")127            128            total_images += processed129            total_bboxes_all += bboxes130    131    print(f"\n{'='*50}")132    print(f"转换完成！")133    print(f"总图片数: {total_images}")134    print(f"总标注数: {total_bboxes_all}")135    print(f"输出文件: {output_path}")136 137if __name__ == "__main__":138    main()