dataset-toolkit 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

dataset_toolkit/__init__.py

@@ -15,9 +15,9 @@ Dataset Toolkit - 计算机视觉数据集处理工具包
     >>> export_to_coco(dataset, "output.json")
 """
 
-__version__ = "0.1.0"
-__author__ = "Your Name"
-__email__ = "your.email@example.com"
+__version__ = "0.1.2"
+__author__ = "wenxiang.han"
+__email__ = "wenxiang.han@anker-in.com"
 
 # 导入核心类和函数，提供简洁的顶层API
 from dataset_toolkit.models import (
@@ -27,7 +27,8 @@ from dataset_toolkit.models import (
 )
 
 from dataset_toolkit.loaders.local_loader import (
-    load_yolo_from_local
+    load_yolo_from_local,
+    load_csv_result_from_local
 )
 
 from dataset_toolkit.processors.merger import (
@@ -42,6 +43,11 @@ from dataset_toolkit.exporters.txt_exporter import (
     export_to_txt
 )
 
+from dataset_toolkit.exporters.yolo_exporter import (
+    export_to_yolo_format,
+    export_to_yolo_and_txt
+)
+
 from dataset_toolkit.utils.coords import (
     yolo_to_absolute_bbox
 )
@@ -62,6 +68,7 @@ __all__ = [
     
     # 加载器
     "load_yolo_from_local",
+    "load_csv_result_from_local",
     
     # 处理器
     "merge_datasets",
@@ -69,6 +76,8 @@ __all__ = [
     # 导出器
     "export_to_coco",
     "export_to_txt",
+    "export_to_yolo_format",
+    "export_to_yolo_and_txt",
     
     # 工具函数
     "yolo_to_absolute_bbox",

dataset_toolkit/exporters/yolo_exporter.py

@@ -0,0 +1,155 @@
+# dataset_toolkit/exporters/yolo_exporter.py
+"""
+导出为 YOLO 格式（完整的 images/ + labels/ 目录结构）
+"""
+import os
+from pathlib import Path
+from typing import Optional
+
+
+def export_to_yolo_format(
+    dataset,
+    output_dir: str,
+    use_symlinks: bool = True,
+    overwrite: bool = False
+):
+    """
+    导出数据集为完整的 YOLO 格式目录结构
+    
+    参数:
+        dataset: Dataset 对象
+        output_dir: 输出目录路径
+        use_symlinks: 是否使用软链接（True）或复制文件（False）
+        overwrite: 是否覆盖已存在的文件
+    
+    输出结构:
+        output_dir/
+        ├── images/
+        │   ├── img1.jpg
+        │   └── img2.jpg
+        └── labels/
+            ├── img1.txt
+            └── img2.txt
+    """
+    output_path = Path(output_dir)
+    images_dir = output_path / 'images'
+    labels_dir = output_path / 'labels'
+    
+    # 创建目录
+    images_dir.mkdir(parents=True, exist_ok=True)
+    labels_dir.mkdir(parents=True, exist_ok=True)
+    
+    print(f"导出 YOLO 格式到: {output_path}")
+    print(f"  使用软链接: {use_symlinks}")
+    
+    success_count = 0
+    error_count = 0
+    
+    for img in dataset.images:
+        try:
+            # 获取图片文件名（不含扩展名）
+            img_path = Path(img.path)
+            img_name = img_path.name
+            stem = img_path.stem
+            
+            # 1. 处理图片（软链接或复制）
+            target_img_path = images_dir / img_name
+            
+            if target_img_path.exists() and not overwrite:
+                # 文件已存在，跳过
+                pass
+            else:
+                if use_symlinks:
+                    # 使用软链接
+                    if target_img_path.exists():
+                        target_img_path.unlink()
+                    target_img_path.symlink_to(img_path.resolve())
+                else:
+                    # 复制文件
+                    import shutil
+                    shutil.copy2(img_path, target_img_path)
+            
+            # 2. 生成标注文件
+            label_path = labels_dir / f"{stem}.txt"
+            
+            with open(label_path, 'w') as f:
+                for ann in img.annotations:
+                    # 内部格式: [x_min, y_min, width, height] (绝对像素值)
+                    # YOLO 格式: class_id x_center y_center width height (归一化)
+                    
+                    x_min, y_min, width, height = ann.bbox
+                    
+                    # 转换为 YOLO 归一化格式
+                    x_center = (x_min + width / 2) / img.width
+                    y_center = (y_min + height / 2) / img.height
+                    norm_width = width / img.width
+                    norm_height = height / img.height
+                    
+                    # 写入：class_id x_center y_center width height
+                    f.write(f"{ann.category_id} {x_center:.6f} {y_center:.6f} {norm_width:.6f} {norm_height:.6f}\n")
+            
+            success_count += 1
+            
+        except Exception as e:
+            print(f"警告: 处理图片失败 {img.path}: {e}")
+            error_count += 1
+            continue
+    
+    print(f"✓ 导出完成:")
+    print(f"  成功: {success_count} 张图片")
+    if error_count > 0:
+        print(f"  失败: {error_count} 张图片")
+    print(f"  图片目录: {images_dir}")
+    print(f"  标注目录: {labels_dir}")
+    
+    return output_path
+
+
+def export_to_yolo_and_txt(
+    dataset,
+    yolo_dir: str,
+    txt_file: str,
+    use_symlinks: bool = True,
+    use_relative_paths: bool = False
+):
+    """
+    导出为 YOLO 格式并生成对应的 txt 列表文件
+    
+    参数:
+        dataset: Dataset 对象
+        yolo_dir: YOLO 格式输出目录
+        txt_file: txt 列表文件路径
+        use_symlinks: 是否使用软链接
+        use_relative_paths: txt 中是否使用相对路径
+    
+    返回:
+        yolo_dir_path: YOLO 目录路径
+    """
+    # 1. 导出为 YOLO 格式
+    yolo_path = export_to_yolo_format(dataset, yolo_dir, use_symlinks=use_symlinks)
+    
+    # 2. 生成 txt 列表文件（指向 YOLO 目录中的 images/）
+    images_dir = yolo_path / 'images'
+    txt_path = Path(txt_file)
+    txt_path.parent.mkdir(parents=True, exist_ok=True)
+    
+    print(f"\n生成 txt 列表: {txt_file}")
+    
+    with open(txt_file, 'w') as f:
+        for img in dataset.images:
+            img_name = Path(img.path).name
+            # 指向 YOLO 目录中的图片（可能是软链接）
+            img_in_yolo = images_dir / img_name
+            
+            if use_relative_paths:
+                # 相对于 txt 文件的路径
+                rel_path = os.path.relpath(img_in_yolo, txt_path.parent)
+                f.write(f"{rel_path}\n")
+            else:
+                # 绝对路径（指向 YOLO images 目录，不要 resolve，保持 YOLO 结构）
+                f.write(f"{str(img_in_yolo.absolute())}\n")
+    
+    print(f"✓ txt 列表已生成: {len(dataset.images)} 行")
+    
+    return yolo_path
+

dataset_toolkit/loaders/local_loader.py

@@ -2,6 +2,8 @@
 from pathlib import Path
 from typing import Dict
 from PIL import Image
+import csv
+import json
 
 # 从我们自己的包中导入模块
 from dataset_toolkit.models import Dataset, ImageAnnotation, Annotation
@@ -61,4 +63,127 @@ def load_yolo_from_local(dataset_path: str, categories: Dict[int, str]) -> Datas
         dataset.images.append(image_annotation)
 
     print(f"加载完成. 共找到 {len(dataset.images)} 张图片.")
+    return dataset
+
+
+def load_csv_result_from_local(dataset_path: str, categories: Dict[int, str] = None) -> Dataset:
+    """
+    从本地文件系统加载包含 result.csv 的数据集。
+    
+    数据集结构：
+    - 根目录下包含 jpg 图片文件
+    - result.csv 文件，格式为：file_id,result_json
+    - result_json 是 JSON 数组，包含检测结果
+      格式: [{"box": [x1, y1, x2, y2], "conf": 0.xx, "class_id": 0, "class_name": "parcel"}]
+    
+    参数:
+        dataset_path: 数据集根目录路径
+        categories: 类别映射字典 {class_id: class_name}，如果为 None 则从数据自动提取
+    """
+    root_path = Path(dataset_path)
+    csv_path = root_path / 'result.csv'
+    
+    if not csv_path.exists():
+        raise FileNotFoundError(f"result.csv 文件不存在: {csv_path}")
+    
+    # 如果没有提供 categories，则使用空字典，稍后从数据中提取
+    if categories is None:
+        categories = {}
+    
+    dataset = Dataset(name=root_path.name, categories=categories)
+    supported_extensions = ['.jpg', '.jpeg', '.png']
+    
+    print(f"开始加载数据集: {root_path.name}...")
+    
+    # 读取 CSV 文件，建立 file_id 到 result_json 的映射
+    # 手动解析以处理 JSON 中的逗号
+    results_dict = {}
+    with open(csv_path, 'r', encoding='utf-8') as f:
+        # 跳过表头
+        header = f.readline().strip()
+        if header != 'file_id,result_json':
+            print(f"警告: CSV 表头格式不匹配，期望 'file_id,result_json'，实际为 '{header}'")
+        
+        # 逐行解析
+        for line_num, line in enumerate(f, start=2):
+            line = line.strip()
+            if not line:
+                continue
+            
+            # 查找第一个逗号作为分隔符，之后的所有内容都是 result_json
+            comma_idx = line.find(',')
+            if comma_idx == -1:
+                print(f"警告: 第 {line_num} 行格式错误，已跳过: {line}")
+                continue
+            
+            file_id = line[:comma_idx].strip()
+            result_json = line[comma_idx + 1:].strip()
+            results_dict[file_id] = result_json
+    
+    print(f"从 result.csv 读取了 {len(results_dict)} 条标注记录.")
+    
+    # 遍历根目录下的所有图片文件
+    image_count = 0
+    for image_path in root_path.iterdir():
+        if not image_path.is_file() or image_path.suffix.lower() not in supported_extensions:
+            continue
+        
+        try:
+            with Image.open(image_path) as img:
+                img_width, img_height = img.size
+        except IOError:
+            print(f"警告: 无法打开图片，已跳过: {image_path}")
+            continue
+        
+        image_annotation = ImageAnnotation(
+            image_id=image_path.name,
+            path=str(image_path.resolve()),
+            width=img_width,
+            height=img_height
+        )
+        
+        # 查找对应的标注结果（文件名不含后缀）
+        file_id = image_path.stem
+        if file_id in results_dict:
+            result_json = results_dict[file_id]
+            
+            # 解析 JSON 格式的检测结果
+            try:
+                detections = json.loads(result_json)
+                
+                # 遍历每个检测框
+                for det in detections:
+                    box = det.get('box', [])
+                    conf = det.get('conf', 1.0)
+                    class_id = det.get('class_id', 0)
+                    class_name = det.get('class_name', 'unknown')
+                    
+                    # 如果 categories 中没有这个类别，自动添加
+                    if class_id not in dataset.categories:
+                        dataset.categories[class_id] = class_name
+                    
+                    # box 格式为 [x1, y1, x2, y2]，需要转换为 [x_min, y_min, width, height]
+                    if len(box) == 4:
+                        x1, y1, x2, y2 = box
+                        x_min = x1
+                        y_min = y1
+                        width = x2 - x1
+                        height = y2 - y1
+                        
+                        annotation = Annotation(
+                            category_id=class_id,
+                            bbox=[x_min, y_min, width, height]
+                        )
+                        image_annotation.annotations.append(annotation)
+                    else:
+                        print(f"警告: 无效的边界框格式，已跳过: {file_id} -> {box}")
+                        
+            except json.JSONDecodeError as e:
+                print(f"警告: 无法解析 JSON，已跳过: {file_id} -> {e}")
+        
+        dataset.images.append(image_annotation)
+        image_count += 1
+    
+    print(f"加载完成. 共找到 {image_count} 张图片, {len(dataset.categories)} 个类别.")
+    print(f"类别映射: {dataset.categories}")
     return dataset

dataset_toolkit/models.py

@@ -8,6 +8,7 @@ class Annotation:
     category_id: int
     # 存储格式为 [x_min, y_min, width, height]，单位是绝对像素值
     bbox: List[float]
+    confidence: float = 1.0  # 检测置信度，默认为 1.0
 
 @dataclass
 class ImageAnnotation:

{dataset_toolkit-0.1.0.dist-info → dataset_toolkit-0.1.2.dist-info}/METADATA

@@ -1,10 +1,10 @@
 Metadata-Version: 2.4
 Name: dataset-toolkit
-Version: 0.1.0
+Version: 0.1.2
 Summary: 一个用于加载、处理和导出计算机视觉数据集的工具包
 Home-page: https://github.com/yourusername/dataset-toolkit
-Author: Your Name
-Author-email: Your Name <your.email@example.com>
+Author: wenxiang.han
+Author-email: "wenxiang.han" <wenxiang.han@anker-in.com>
 License: MIT
 Project-URL: Homepage, https://github.com/yourusername/dataset-toolkit
 Project-URL: Documentation, https://dataset-toolkit.readthedocs.io

{dataset_toolkit-0.1.0.dist-info → dataset_toolkit-0.1.2.dist-info}/RECORD

@@ -1,17 +1,18 @@
-dataset_toolkit/__init__.py,sha256=yMm1ajpItXWlKdiqEmY3kRXDI9F0Voreg4hEH0xxM1s,1604
-dataset_toolkit/models.py,sha256=uVtTbVYdHMECPL_waDhEebLvL_VwqSEm9XFC5QYIB10,767
+dataset_toolkit/__init__.py,sha256=3KUOm1r2ldAGf6V0zPbx69QaoJvC9fsuasAvstcN4Vs,1846
+dataset_toolkit/models.py,sha256=9HD2lAOPuEytFb1qRejODLJAD-uKHc8Ya1n9nbGhRpg,830
 dataset_toolkit/pipeline.py,sha256=iBJD7SemEVFTwzHxRQrjpUIQQcVdPSZnD4sB_y56Md0,5697
 dataset_toolkit/exporters/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dataset_toolkit/exporters/coco_exporter.py,sha256=l5sfj7rOcvcMC0-4LNOEJ4PeklGQORDflU_um5GGnxA,2120
 dataset_toolkit/exporters/txt_exporter.py,sha256=9nTWs6M89MdKJhlODtmfzeZqWkliXac9NMWPgVUrE7c,1246
+dataset_toolkit/exporters/yolo_exporter.py,sha256=g0jaY6cg7I4ZfxXkE_vmK87OS1-DP6dXEjsP2iVP9D4,5144
 dataset_toolkit/loaders/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-dataset_toolkit/loaders/local_loader.py,sha256=Wy_hXY2B-SDxAmJGBYQpqBUe3cjz-k_McYhYf7cLgCk,2501
+dataset_toolkit/loaders/local_loader.py,sha256=SCOYG5pursEIL_m3QYGcm-2skXoapiOA4yhqqa2wrDM,7468
 dataset_toolkit/processors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dataset_toolkit/processors/merger.py,sha256=h8qQNgSmkPrhoQ3QiWEyIl11CmmjT5K1-8TzNb7_jbk,2834
 dataset_toolkit/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dataset_toolkit/utils/coords.py,sha256=GtTQz2gFyFQfXhKfecI8tzqWFjraJY6Xo85-kRXYAYc,614
-dataset_toolkit-0.1.0.dist-info/licenses/LICENSE,sha256=8_up1FX6vk2DRcusQEZ4pWJGkgkjvEkD14xB1hdLe3c,1067
-dataset_toolkit-0.1.0.dist-info/METADATA,sha256=yjqPr_Wjioiw5v7AOkqduI5B_Y6oyBbKrTpJGIKVIWw,7225
-dataset_toolkit-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-dataset_toolkit-0.1.0.dist-info/top_level.txt,sha256=B4D5vMLjUNJBZDdL7Utc0FYIfYoWbzyIGBMVYaeMd3U,16
-dataset_toolkit-0.1.0.dist-info/RECORD,,
+dataset_toolkit-0.1.2.dist-info/licenses/LICENSE,sha256=8_up1FX6vk2DRcusQEZ4pWJGkgkjvEkD14xB1hdLe3c,1067
+dataset_toolkit-0.1.2.dist-info/METADATA,sha256=qH0LX-6NPdZEd0htGIxuIMHueDmHxRsZmWbrGnbpEtA,7236
+dataset_toolkit-0.1.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+dataset_toolkit-0.1.2.dist-info/top_level.txt,sha256=B4D5vMLjUNJBZDdL7Utc0FYIfYoWbzyIGBMVYaeMd3U,16
+dataset_toolkit-0.1.2.dist-info/RECORD,,