You've already forked DataMate
Jerry Yan
3c3ca130b3
- 新增异步函数 _read_file_content 用于安全读取文本文件内容 - 实现在导出时包含文本文件的实际内容数据 - 扩展 CSV 导出格式支持多种标注类型标签提取 - 添加对矩形标签、多边形标签、画笔标签等多种标注类型的支持 - 更新 COCO 格式导出文档说明bbox坐标转换注意事项
484 lines
18 KiB
Python
484 lines
18 KiB
Python
"""
|
|
标注数据导出服务
|
|
|
|
支持的导出格式：
|
|
- JSON: Label Studio 原生 JSON 格式
|
|
- JSONL: JSON Lines 格式（每行一条记录）
|
|
- CSV: CSV 表格格式
|
|
- COCO: COCO 目标检测格式（适用于图像标注）
|
|
- YOLO: YOLO 格式（适用于图像标注）
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import csv
|
|
import io
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import zipfile
|
|
from datetime import datetime
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
from fastapi import HTTPException
|
|
from sqlalchemy import func, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.logging import get_logger
|
|
from app.db.models import AnnotationResult, Dataset, DatasetFiles, LabelingProject
|
|
|
|
|
|
async def _read_file_content(file_path: str, max_size: int = 10 * 1024 * 1024) -> Optional[str]:
|
|
"""读取文件内容，仅适用于文本文件
|
|
|
|
Args:
|
|
file_path: 文件路径
|
|
max_size: 最大读取字节数（默认10MB）
|
|
|
|
Returns:
|
|
文件内容字符串，如果读取失败返回 None
|
|
"""
|
|
try:
|
|
# 检查文件是否存在且大小在限制内
|
|
if not os.path.exists(file_path):
|
|
return None
|
|
|
|
file_size = os.path.getsize(file_path)
|
|
if file_size > max_size:
|
|
return f"[File too large: {file_size} bytes]"
|
|
|
|
# 尝试以文本方式读取
|
|
with open(file_path, 'r', encoding='utf-8', errors='replace') as f:
|
|
return f.read()
|
|
except Exception:
|
|
return None
|
|
|
|
from ..schema.export import (
|
|
AnnotationExportItem,
|
|
COCOExportFormat,
|
|
ExportAnnotationsRequest,
|
|
ExportAnnotationsResponse,
|
|
ExportFormat,
|
|
)
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class AnnotationExportService:
|
|
"""标注数据导出服务"""
|
|
|
|
def __init__(self, db: AsyncSession):
|
|
self.db = db
|
|
|
|
async def get_export_stats(self, project_id: str) -> ExportAnnotationsResponse:
|
|
"""获取导出统计信息"""
|
|
project = await self._get_project_or_404(project_id)
|
|
logger.info(f"Export stats for project: id={project_id}, dataset_id={project.dataset_id}, name={project.name}")
|
|
|
|
# 获取总文件数（只统计 ACTIVE 状态的文件）
|
|
total_result = await self.db.execute(
|
|
select(func.count()).select_from(DatasetFiles).where(
|
|
DatasetFiles.dataset_id == project.dataset_id,
|
|
DatasetFiles.status == "ACTIVE",
|
|
)
|
|
)
|
|
total_files = int(total_result.scalar() or 0)
|
|
logger.info(f"Total files (ACTIVE): {total_files} for dataset_id={project.dataset_id}")
|
|
|
|
# 获取已标注文件数（统计不同的 file_id 数量）
|
|
annotated_result = await self.db.execute(
|
|
select(func.count(func.distinct(AnnotationResult.file_id))).where(
|
|
AnnotationResult.project_id == project_id
|
|
)
|
|
)
|
|
annotated_files = int(annotated_result.scalar() or 0)
|
|
logger.info(f"Annotated files: {annotated_files} for project_id={project_id}")
|
|
|
|
return ExportAnnotationsResponse(
|
|
project_id=project_id,
|
|
project_name=project.name,
|
|
total_files=total_files,
|
|
annotated_files=annotated_files,
|
|
export_format="json",
|
|
)
|
|
|
|
async def export_annotations(
|
|
self,
|
|
project_id: str,
|
|
request: ExportAnnotationsRequest,
|
|
) -> Tuple[bytes, str, str]:
|
|
"""
|
|
导出标注数据
|
|
|
|
返回: (文件内容bytes, 文件名, content_type)
|
|
"""
|
|
project = await self._get_project_or_404(project_id)
|
|
|
|
# 获取标注数据
|
|
items = await self._fetch_annotation_data(
|
|
project_id=project_id,
|
|
dataset_id=project.dataset_id,
|
|
only_annotated=request.only_annotated,
|
|
include_data=request.include_data,
|
|
)
|
|
|
|
# 根据格式导出
|
|
format_type = ExportFormat(request.format) if isinstance(request.format, str) else request.format
|
|
|
|
if format_type == ExportFormat.JSON:
|
|
return self._export_json(items, project.name)
|
|
elif format_type == ExportFormat.JSONL:
|
|
return self._export_jsonl(items, project.name)
|
|
elif format_type == ExportFormat.CSV:
|
|
return self._export_csv(items, project.name)
|
|
elif format_type == ExportFormat.COCO:
|
|
return self._export_coco(items, project.name)
|
|
elif format_type == ExportFormat.YOLO:
|
|
return self._export_yolo(items, project.name)
|
|
else:
|
|
raise HTTPException(status_code=400, detail=f"不支持的导出格式: {request.format}")
|
|
|
|
async def _get_project_or_404(self, project_id: str) -> LabelingProject:
|
|
"""获取标注项目，不存在则抛出 404"""
|
|
result = await self.db.execute(
|
|
select(LabelingProject).where(
|
|
LabelingProject.id == project_id,
|
|
LabelingProject.deleted_at.is_(None),
|
|
)
|
|
)
|
|
project = result.scalar_one_or_none()
|
|
if not project:
|
|
raise HTTPException(status_code=404, detail=f"标注项目不存在: {project_id}")
|
|
return project
|
|
|
|
async def _fetch_annotation_data(
|
|
self,
|
|
project_id: str,
|
|
dataset_id: str,
|
|
only_annotated: bool = True,
|
|
include_data: bool = False,
|
|
) -> List[AnnotationExportItem]:
|
|
"""获取标注数据列表"""
|
|
items: List[AnnotationExportItem] = []
|
|
|
|
if only_annotated:
|
|
# 只获取已标注的数据
|
|
result = await self.db.execute(
|
|
select(AnnotationResult, DatasetFiles)
|
|
.join(DatasetFiles, AnnotationResult.file_id == DatasetFiles.id)
|
|
.where(AnnotationResult.project_id == project_id)
|
|
.order_by(AnnotationResult.updated_at.desc())
|
|
)
|
|
rows = result.all()
|
|
|
|
for ann, file in rows:
|
|
annotation_data = ann.annotation or {}
|
|
# 获取文件内容（如果是文本文件且用户要求包含数据）
|
|
file_content = None
|
|
if include_data:
|
|
file_path = getattr(file, "file_path", "")
|
|
file_content = await _read_file_content(file_path)
|
|
|
|
items.append(
|
|
AnnotationExportItem(
|
|
file_id=str(file.id),
|
|
file_name=str(getattr(file, "file_name", "")),
|
|
data={"text": file_content} if include_data else None,
|
|
annotations=[annotation_data] if annotation_data else [],
|
|
created_at=ann.created_at,
|
|
updated_at=ann.updated_at,
|
|
)
|
|
)
|
|
else:
|
|
# 获取所有文件，包括未标注的（只获取 ACTIVE 状态的文件）
|
|
files_result = await self.db.execute(
|
|
select(DatasetFiles).where(
|
|
DatasetFiles.dataset_id == dataset_id,
|
|
DatasetFiles.status == "ACTIVE",
|
|
)
|
|
)
|
|
files = files_result.scalars().all()
|
|
|
|
# 获取已有的标注
|
|
ann_result = await self.db.execute(
|
|
select(AnnotationResult).where(AnnotationResult.project_id == project_id)
|
|
)
|
|
annotations = {str(a.file_id): a for a in ann_result.scalars().all()}
|
|
|
|
for file in files:
|
|
file_id = str(file.id)
|
|
ann = annotations.get(file_id)
|
|
annotation_data = ann.annotation if ann else {}
|
|
|
|
# 获取文件内容（如果是文本文件且用户要求包含数据）
|
|
file_content = None
|
|
if include_data:
|
|
file_path = getattr(file, "file_path", "")
|
|
file_content = await _read_file_content(file_path)
|
|
|
|
items.append(
|
|
AnnotationExportItem(
|
|
file_id=file_id,
|
|
file_name=str(getattr(file, "file_name", "")),
|
|
data={"text": file_content} if include_data else None,
|
|
annotations=[annotation_data] if annotation_data else [],
|
|
created_at=ann.created_at if ann else None,
|
|
updated_at=ann.updated_at if ann else None,
|
|
)
|
|
)
|
|
|
|
return items
|
|
|
|
def _export_json(
|
|
self, items: List[AnnotationExportItem], project_name: str
|
|
) -> Tuple[bytes, str, str]:
|
|
"""导出为 JSON 格式"""
|
|
export_data = {
|
|
"project_name": project_name,
|
|
"export_time": datetime.utcnow().isoformat() + "Z",
|
|
"total_items": len(items),
|
|
"annotations": [
|
|
{
|
|
"file_id": item.file_id,
|
|
"file_name": item.file_name,
|
|
"data": item.data,
|
|
"annotations": item.annotations,
|
|
"created_at": item.created_at.isoformat() if item.created_at else None,
|
|
"updated_at": item.updated_at.isoformat() if item.updated_at else None,
|
|
}
|
|
for item in items
|
|
],
|
|
}
|
|
|
|
content = json.dumps(export_data, ensure_ascii=False, indent=2).encode("utf-8")
|
|
filename = f"{project_name}_annotations.json"
|
|
return content, filename, "application/json"
|
|
|
|
def _export_jsonl(
|
|
self, items: List[AnnotationExportItem], project_name: str
|
|
) -> Tuple[bytes, str, str]:
|
|
"""导出为 JSON Lines 格式"""
|
|
lines = []
|
|
for item in items:
|
|
record = {
|
|
"file_id": item.file_id,
|
|
"file_name": item.file_name,
|
|
"data": item.data,
|
|
"annotations": item.annotations,
|
|
"created_at": item.created_at.isoformat() if item.created_at else None,
|
|
"updated_at": item.updated_at.isoformat() if item.updated_at else None,
|
|
}
|
|
lines.append(json.dumps(record, ensure_ascii=False))
|
|
|
|
content = "\n".join(lines).encode("utf-8")
|
|
filename = f"{project_name}_annotations.jsonl"
|
|
return content, filename, "application/x-ndjson"
|
|
|
|
def _export_csv(
|
|
self, items: List[AnnotationExportItem], project_name: str
|
|
) -> Tuple[bytes, str, str]:
|
|
"""导出为 CSV 格式"""
|
|
output = io.StringIO()
|
|
|
|
# CSV 表头
|
|
fieldnames = [
|
|
"file_id",
|
|
"file_name",
|
|
"annotation_result",
|
|
"labels",
|
|
"created_at",
|
|
"updated_at",
|
|
]
|
|
|
|
writer = csv.DictWriter(output, fieldnames=fieldnames)
|
|
writer.writeheader()
|
|
|
|
for item in items:
|
|
# 提取标签信息（支持多种标注类型）
|
|
labels = []
|
|
for ann in item.annotations:
|
|
results = ann.get("result", [])
|
|
for r in results:
|
|
value = r.get("value", {})
|
|
label_type = r.get("type", "")
|
|
|
|
# 提取不同类型的标签值
|
|
if "choices" in value:
|
|
labels.extend(value["choices"])
|
|
elif "text" in value:
|
|
labels.append(value["text"])
|
|
elif "labels" in value:
|
|
labels.extend(value["labels"])
|
|
elif "rectanglelabels" in value:
|
|
labels.extend(value["rectanglelabels"])
|
|
elif "polygonlabels" in value:
|
|
labels.extend(value["polygonlabels"])
|
|
elif "brushlabels" in value:
|
|
labels.extend(value["brushlabels"])
|
|
elif "hypertextlabels" in value:
|
|
labels.extend(value["hypertextlabels"])
|
|
elif "timeserieslabels" in value:
|
|
labels.extend(value["timeserieslabels"])
|
|
elif "transcription" in value:
|
|
labels.append(value["transcription"])
|
|
|
|
writer.writerow({
|
|
"file_id": item.file_id,
|
|
"file_name": item.file_name,
|
|
"annotation_result": json.dumps(item.annotations, ensure_ascii=False),
|
|
"labels": "|".join(labels),
|
|
"created_at": item.created_at.isoformat() if item.created_at else "",
|
|
"updated_at": item.updated_at.isoformat() if item.updated_at else "",
|
|
})
|
|
|
|
content = output.getvalue().encode("utf-8-sig") # BOM for Excel compatibility
|
|
filename = f"{project_name}_annotations.csv"
|
|
return content, filename, "text/csv"
|
|
|
|
def _export_coco(
|
|
self, items: List[AnnotationExportItem], project_name: str
|
|
) -> Tuple[bytes, str, str]:
|
|
"""导出为 COCO 格式（适用于目标检测标注）
|
|
|
|
注意：当前实现中图片宽高被设置为0，因为需要读取实际图片文件获取尺寸。
|
|
bbox 坐标使用 Label Studio 的百分比值（0-100），使用时需要转换为像素坐标。
|
|
"""
|
|
coco_format = COCOExportFormat(
|
|
info={
|
|
"description": f"Exported from DataMate project: {project_name}",
|
|
"version": "1.0",
|
|
"year": datetime.utcnow().year,
|
|
"date_created": datetime.utcnow().isoformat(),
|
|
},
|
|
licenses=[{"id": 1, "name": "Unknown", "url": ""}],
|
|
images=[],
|
|
annotations=[],
|
|
categories=[],
|
|
)
|
|
|
|
category_map: Dict[str, int] = {}
|
|
annotation_id = 1
|
|
|
|
for idx, item in enumerate(items):
|
|
image_id = idx + 1
|
|
|
|
# 添加图片信息
|
|
coco_format.images.append({
|
|
"id": image_id,
|
|
"file_name": item.file_name,
|
|
"width": 0, # 需要实际图片尺寸
|
|
"height": 0,
|
|
})
|
|
|
|
# 处理标注
|
|
for ann in item.annotations:
|
|
results = ann.get("result", [])
|
|
for r in results:
|
|
# 处理矩形框标注 (rectanglelabels)
|
|
if r.get("type") == "rectanglelabels":
|
|
value = r.get("value", {})
|
|
labels = value.get("rectanglelabels", [])
|
|
|
|
for label in labels:
|
|
if label not in category_map:
|
|
category_map[label] = len(category_map) + 1
|
|
coco_format.categories.append({
|
|
"id": category_map[label],
|
|
"name": label,
|
|
"supercategory": "",
|
|
})
|
|
|
|
# 转换坐标（Label Studio 使用百分比）
|
|
x = value.get("x", 0)
|
|
y = value.get("y", 0)
|
|
width = value.get("width", 0)
|
|
height = value.get("height", 0)
|
|
|
|
coco_format.annotations.append({
|
|
"id": annotation_id,
|
|
"image_id": image_id,
|
|
"category_id": category_map[label],
|
|
"bbox": [x, y, width, height],
|
|
"area": width * height,
|
|
"iscrowd": 0,
|
|
})
|
|
annotation_id += 1
|
|
|
|
content = json.dumps(coco_format.model_dump(), ensure_ascii=False, indent=2).encode("utf-8")
|
|
filename = f"{project_name}_coco.json"
|
|
return content, filename, "application/json"
|
|
|
|
def _export_yolo(
|
|
self, items: List[AnnotationExportItem], project_name: str
|
|
) -> Tuple[bytes, str, str]:
|
|
"""导出为 YOLO 格式（ZIP 包含 txt 标注文件和 classes.txt）"""
|
|
|
|
# 创建临时目录
|
|
tmp_fd, tmp_path = tempfile.mkstemp(suffix=".zip")
|
|
os.close(tmp_fd)
|
|
|
|
category_set: set = set()
|
|
txt_files: Dict[str, str] = {}
|
|
|
|
for item in items:
|
|
lines = []
|
|
|
|
for ann in item.annotations:
|
|
results = ann.get("result", [])
|
|
for r in results:
|
|
# 处理矩形框标注
|
|
if r.get("type") == "rectanglelabels":
|
|
value = r.get("value", {})
|
|
labels = value.get("rectanglelabels", [])
|
|
|
|
for label in labels:
|
|
category_set.add(label)
|
|
|
|
# YOLO 格式：class_id x_center y_center width height (归一化 0-1)
|
|
x = value.get("x", 0) / 100
|
|
y = value.get("y", 0) / 100
|
|
w = value.get("width", 0) / 100
|
|
h = value.get("height", 0) / 100
|
|
|
|
x_center = x + w / 2
|
|
y_center = y + h / 2
|
|
|
|
lines.append(f"{label} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}")
|
|
|
|
if lines:
|
|
# 生成对应的 txt 文件名
|
|
base_name = os.path.splitext(item.file_name)[0]
|
|
txt_files[f"{base_name}.txt"] = "\n".join(lines)
|
|
|
|
# 生成类别列表
|
|
categories = sorted(category_set)
|
|
category_map = {cat: idx for idx, cat in enumerate(categories)}
|
|
|
|
# 更新 txt 文件中的类别索引
|
|
for filename, content in txt_files.items():
|
|
updated_lines = []
|
|
for line in content.split("\n"):
|
|
parts = line.split(" ", 1)
|
|
if len(parts) == 2:
|
|
label, coords = parts
|
|
class_id = category_map.get(label, 0)
|
|
updated_lines.append(f"{class_id} {coords}")
|
|
txt_files[filename] = "\n".join(updated_lines)
|
|
|
|
# 创建 ZIP 文件
|
|
with zipfile.ZipFile(tmp_path, "w", zipfile.ZIP_DEFLATED) as zf:
|
|
# 写入 classes.txt
|
|
zf.writestr("classes.txt", "\n".join(categories))
|
|
|
|
# 写入标注文件
|
|
for filename, content in txt_files.items():
|
|
zf.writestr(f"labels/{filename}", content)
|
|
|
|
with open(tmp_path, "rb") as f:
|
|
content = f.read()
|
|
|
|
os.unlink(tmp_path)
|
|
|
|
filename = f"{project_name}_yolo.zip"
|
|
return content, filename, "application/zip"
|