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

dataset_toolkit/__init__.py

@@ -15,7 +15,7 @@ Dataset Toolkit - 计算机视觉数据集处理工具包
     >>> export_to_coco(dataset, "output.json")
 """
 
-__version__ = "0.1.1"
+__version__ = "0.1.2"
 __author__ = "wenxiang.han"
 __email__ = "wenxiang.han@anker-in.com"
 
@@ -43,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
 )
@@ -71,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-0.1.1.dist-info → dataset_toolkit-0.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataset-toolkit
-Version: 0.1.1
+Version: 0.1.2
 Summary: 一个用于加载、处理和导出计算机视觉数据集的工具包
 Home-page: https://github.com/yourusername/dataset-toolkit
 Author: wenxiang.han

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

@@ -1,17 +1,18 @@
-dataset_toolkit/__init__.py,sha256=BhgTMzT5onSjI-Sd2bFSQGYyo9GwcLZUyowlyx7lMyU,1676
+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=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.1.dist-info/licenses/LICENSE,sha256=8_up1FX6vk2DRcusQEZ4pWJGkgkjvEkD14xB1hdLe3c,1067
-dataset_toolkit-0.1.1.dist-info/METADATA,sha256=l3COSL22yVvDDZL_c_N5uJNjAPpKE0o2BasMso_Ntss,7236
-dataset_toolkit-0.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-dataset_toolkit-0.1.1.dist-info/top_level.txt,sha256=B4D5vMLjUNJBZDdL7Utc0FYIfYoWbzyIGBMVYaeMd3U,16
-dataset_toolkit-0.1.1.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,,