dataset-toolkit 0.1.1__tar.gz → 0.2.0__tar.gz

{dataset_toolkit-0.1.1/dataset_toolkit.egg-info → dataset_toolkit-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataset-toolkit
-Version: 0.1.1
+Version: 0.2.0
 Summary: 一个用于加载、处理和导出计算机视觉数据集的工具包
 Home-page: https://github.com/yourusername/dataset-toolkit
 Author: wenxiang.han
@@ -42,6 +42,7 @@ Dynamic: requires-python
 - 📤 **灵活导出**：导出为 COCO JSON、TXT 等多种格式
 - 🛠️ **工具函数**：提供坐标转换等实用工具
 - 📦 **标准化数据模型**：统一的内部数据表示，方便扩展
+- 📊 **模型评估**：完整的目标检测模型评估系统（v0.2.0+）
 
 ## 📦 安装
 
@@ -121,6 +122,43 @@ result = (pipeline
     .execute())
 ```
 
+### 模型评估（v0.2.0+）
+
+```python
+from dataset_toolkit import (
+    load_yolo_from_local,
+    load_predictions_from_streamlined,
+    Evaluator
+)
+
+# 1. 加载GT和预测结果
+gt_dataset = load_yolo_from_local("/data/test/labels", {0: 'parcel'})
+pred_dataset = load_predictions_from_streamlined(
+    "/results/predictions",
+    categories={0: 'parcel'},
+    image_dir="/data/test/images"
+)
+
+# 2. 创建评估器
+evaluator = Evaluator(
+    positive_gt=gt_dataset,
+    positive_pred=pred_dataset,
+    iou_threshold=0.5
+)
+
+# 3. 计算指标
+metrics = evaluator.calculate_metrics(confidence_threshold=0.5)
+print(f"Precision: {metrics['precision']:.4f}")
+print(f"Recall: {metrics['recall']:.4f}")
+print(f"F1-Score: {metrics['f1']:.4f}")
+
+# 4. 寻找最优阈值
+optimal = evaluator.find_optimal_threshold(metric='f1')
+print(f"最优阈值: {optimal['optimal_threshold']}")
+```
+
+详细文档请参考 [EVALUATION_GUIDE.md](EVALUATION_GUIDE.md)
+
 ## 📚 API 文档
 
 ### 数据加载器

{dataset_toolkit-0.1.1 → dataset_toolkit-0.2.0}/README.md

@@ -9,6 +9,7 @@
 - 📤 **灵活导出**：导出为 COCO JSON、TXT 等多种格式
 - 🛠️ **工具函数**：提供坐标转换等实用工具
 - 📦 **标准化数据模型**：统一的内部数据表示，方便扩展
+- 📊 **模型评估**：完整的目标检测模型评估系统（v0.2.0+）
 
 ## 📦 安装
 
@@ -88,6 +89,43 @@ result = (pipeline
     .execute())
 ```
 
+### 模型评估（v0.2.0+）
+
+```python
+from dataset_toolkit import (
+    load_yolo_from_local,
+    load_predictions_from_streamlined,
+    Evaluator
+)
+
+# 1. 加载GT和预测结果
+gt_dataset = load_yolo_from_local("/data/test/labels", {0: 'parcel'})
+pred_dataset = load_predictions_from_streamlined(
+    "/results/predictions",
+    categories={0: 'parcel'},
+    image_dir="/data/test/images"
+)
+
+# 2. 创建评估器
+evaluator = Evaluator(
+    positive_gt=gt_dataset,
+    positive_pred=pred_dataset,
+    iou_threshold=0.5
+)
+
+# 3. 计算指标
+metrics = evaluator.calculate_metrics(confidence_threshold=0.5)
+print(f"Precision: {metrics['precision']:.4f}")
+print(f"Recall: {metrics['recall']:.4f}")
+print(f"F1-Score: {metrics['f1']:.4f}")
+
+# 4. 寻找最优阈值
+optimal = evaluator.find_optimal_threshold(metric='f1')
+print(f"最优阈值: {optimal['optimal_threshold']}")
+```
+
+详细文档请参考 [EVALUATION_GUIDE.md](EVALUATION_GUIDE.md)
+
 ## 📚 API 文档
 
 ### 数据加载器

{dataset_toolkit-0.1.1 → dataset_toolkit-0.2.0}/dataset_toolkit/__init__.py

@@ -15,7 +15,7 @@ Dataset Toolkit - 计算机视觉数据集处理工具包
     >>> export_to_coco(dataset, "output.json")
 """
 
-__version__ = "0.1.1"
+__version__ = "0.2.0"
 __author__ = "wenxiang.han"
 __email__ = "wenxiang.han@anker-in.com"
 
@@ -28,13 +28,18 @@ from dataset_toolkit.models import (
 
 from dataset_toolkit.loaders.local_loader import (
     load_yolo_from_local,
-    load_csv_result_from_local
+    load_csv_result_from_local,
+    load_predictions_from_streamlined
 )
 
 from dataset_toolkit.processors.merger import (
     merge_datasets
 )
 
+from dataset_toolkit.processors.evaluator import (
+    Evaluator
+)
+
 from dataset_toolkit.exporters.coco_exporter import (
     export_to_coco
 )
@@ -43,6 +48,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
 )
@@ -64,13 +74,17 @@ __all__ = [
     # 加载器
     "load_yolo_from_local",
     "load_csv_result_from_local",
+    "load_predictions_from_streamlined",
     
     # 处理器
     "merge_datasets",
+    "Evaluator",
     
     # 导出器
     "export_to_coco",
     "export_to_txt",
+    "export_to_yolo_format",
+    "export_to_yolo_and_txt",
     
     # 工具函数
     "yolo_to_absolute_bbox",

dataset_toolkit-0.2.0/dataset_toolkit/exporters/yolo_exporter.py

@@ -0,0 +1,157 @@
+# 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:
+                # 绝对路径（规范化但不解析软链接）
+                # 使用 os.path.normpath 规范化路径，去除 .. 等
+                normalized_path = os.path.normpath(str(img_in_yolo.absolute()))
+                f.write(f"{normalized_path}\n")
+    
+    print(f"✓ txt 列表已生成: {len(dataset.images)} 行")
+    
+    return yolo_path
+

{dataset_toolkit-0.1.1 → dataset_toolkit-0.2.0}/dataset_toolkit/loaders/local_loader.py

@@ -186,4 +186,149 @@ def load_csv_result_from_local(dataset_path: str, categories: Dict[int, str] = N
     
     print(f"加载完成. 共找到 {image_count} 张图片, {len(dataset.categories)} 个类别.")
     print(f"类别映射: {dataset.categories}")
+    return dataset
+
+
+def load_predictions_from_streamlined(
+    predictions_dir: str, 
+    categories: Dict[int, str],
+    image_dir: str = None
+) -> Dataset:
+    """
+    从streamlined推理结果目录加载预测数据集。
+    
+    预测文件格式（每行一个检测）：
+        class_id,confidence,center_x,center_y,width,height
+        例如: 0,0.934679,354.00,388.00,274.00,102.00
+    
+    参数:
+        predictions_dir: 预测结果txt文件所在目录
+        categories: 类别映射字典 {class_id: class_name}
+        image_dir: 图像目录（可选，用于读取图像尺寸）
+                  如果不提供，将尝试从预测文件同级目录查找
+    
+    返回:
+        Dataset: 预测数据集对象，dataset_type='pred'
+    """
+    pred_path = Path(predictions_dir)
+    
+    if not pred_path.is_dir():
+        raise FileNotFoundError(f"预测结果目录不存在: {pred_path}")
+    
+    # 尝试自动查找图像目录
+    if image_dir is None:
+        # 尝试常见的图像目录位置
+        possible_image_dirs = [
+            pred_path.parent / 'images',
+            pred_path.parent.parent / 'images',
+        ]
+        for possible_dir in possible_image_dirs:
+            if possible_dir.is_dir():
+                image_dir = str(possible_dir)
+                print(f"自动找到图像目录: {image_dir}")
+                break
+    
+    dataset = Dataset(
+        name=pred_path.name, 
+        categories=categories,
+        dataset_type="pred"
+    )
+    
+    supported_extensions = ['.jpg', '.jpeg', '.png']
+    txt_files = list(pred_path.glob('*.txt'))
+    
+    print(f"开始加载预测结果: {pred_path.name}...")
+    print(f"找到 {len(txt_files)} 个预测文件")
+    
+    loaded_count = 0
+    skipped_count = 0
+    
+    for txt_file in txt_files:
+        # 预测文件名对应的图像文件名（假设同名）
+        image_base_name = txt_file.stem
+        
+        # 尝试查找对应的图像文件
+        image_path = None
+        img_width, img_height = None, None
+        
+        if image_dir:
+            image_dir_path = Path(image_dir)
+            for ext in supported_extensions:
+                potential_image = image_dir_path / (image_base_name + ext)
+                if potential_image.exists():
+                    image_path = str(potential_image.resolve())
+                    try:
+                        with Image.open(potential_image) as img:
+                            img_width, img_height = img.size
+                    except IOError:
+                        print(f"警告: 无法打开图片 {potential_image}")
+                    break
+        
+        # 如果没有找到图像，使用默认值
+        if image_path is None:
+            # 假设一个默认的图像路径和尺寸
+            image_path = f"unknown/{image_base_name}.jpg"
+            img_width, img_height = 640, 640  # 默认尺寸
+            if image_dir:
+                skipped_count += 1
+        
+        # 创建图像标注对象
+        image_annotation = ImageAnnotation(
+            image_id=image_base_name + '.jpg',
+            path=image_path,
+            width=img_width,
+            height=img_height
+        )
+        
+        # 读取预测结果
+        try:
+            with open(txt_file, 'r') as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    
+                    # 解析格式: class_id,confidence,center_x,center_y,width,height
+                    parts = line.split(',')
+                    if len(parts) != 6:
+                        print(f"警告: 格式错误，已跳过: {txt_file} -> '{line}'")
+                        continue
+                    
+                    try:
+                        class_id = int(parts[0])
+                        confidence = float(parts[1])
+                        center_x = float(parts[2])
+                        center_y = float(parts[3])
+                        width = float(parts[4])
+                        height = float(parts[5])
+                        
+                        # 转换为 [x_min, y_min, width, height] 格式
+                        x_min = center_x - width / 2
+                        y_min = center_y - height / 2
+                        
+                        annotation = Annotation(
+                            category_id=class_id,
+                            bbox=[x_min, y_min, width, height],
+                            confidence=confidence
+                        )
+                        image_annotation.annotations.append(annotation)
+                        
+                    except (ValueError, IndexError) as e:
+                        print(f"警告: 解析错误，已跳过: {txt_file} -> '{line}' ({e})")
+                        continue
+        
+        except Exception as e:
+            print(f"警告: 读取文件失败，已跳过: {txt_file} ({e})")
+            continue
+        
+        dataset.images.append(image_annotation)
+        loaded_count += 1
+    
+    print(f"加载完成. 成功加载 {loaded_count} 个预测文件")
+    if skipped_count > 0:
+        print(f"警告: {skipped_count} 个文件未找到对应图像，使用默认尺寸")
+    
+    total_detections = sum(len(img.annotations) for img in dataset.images)
+    print(f"总检测数: {total_detections}")
+    
     return dataset

{dataset_toolkit-0.1.1 → dataset_toolkit-0.2.0}/dataset_toolkit/models.py

@@ -24,4 +24,6 @@ class Dataset:
     """代表一个完整的数据集对象，作为系统内部的标准化表示."""
     name: str
     images: List[ImageAnnotation] = field(default_factory=list)
-    categories: Dict[int, str] = field(default_factory=dict)
+    categories: Dict[int, str] = field(default_factory=dict)
+    dataset_type: str = "train"  # 'train', 'gt', 'pred'
+    metadata: Dict = field(default_factory=dict)  # 存储描述性信息，不包含处理参数

dataset_toolkit-0.2.0/dataset_toolkit/processors/__init__.py

@@ -0,0 +1,9 @@
+# dataset_toolkit/processors/__init__.py
+from .merger import merge_datasets
+from .evaluator import Evaluator
+
+__all__ = [
+    'merge_datasets',
+    'Evaluator',
+]
+