dataset-toolkit 0.1.0__py3-none-any.whl → 0.1.1__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.1"
+__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 (
@@ -62,6 +63,7 @@ __all__ = [
     
     # 加载器
     "load_yolo_from_local",
+    "load_csv_result_from_local",
     
     # 处理器
     "merge_datasets",

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.1.dist-info}/METADATA

@@ -1,10 +1,10 @@
 Metadata-Version: 2.4
 Name: dataset-toolkit
-Version: 0.1.0
+Version: 0.1.1
 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.1.dist-info}/RECORD

@@ -1,17 +1,17 @@
-dataset_toolkit/__init__.py,sha256=yMm1ajpItXWlKdiqEmY3kRXDI9F0Voreg4hEH0xxM1s,1604
-dataset_toolkit/models.py,sha256=uVtTbVYdHMECPL_waDhEebLvL_VwqSEm9XFC5QYIB10,767
+dataset_toolkit/__init__.py,sha256=BhgTMzT5onSjI-Sd2bFSQGYyo9GwcLZUyowlyx7lMyU,1676
+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/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.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,,