501 lines
17 KiB
Python
501 lines
17 KiB
Python
|
from datetime import datetime
|
||
|
|
import concurrent.futures
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import threading
|
||
|
|
import time
|
||
|
|
import uuid
|
||
|
|
from typing import Any, Dict, List, Optional
|
||
|
|
|
||
|
|
import requests
|
||
|
|
from django.conf import settings
|
||
|
|
from django.db import connection, close_old_connections
|
||
|
|
from django.db.models import Case, Count, IntegerField, Sum, When
|
||
|
|
from django.shortcuts import get_object_or_404
|
||
|
|
from rest_framework import status
|
||
|
|
from rest_framework.decorators import api_view, permission_classes
|
||
|
|
from rest_framework.pagination import PageNumberPagination
|
||
|
|
from rest_framework.permissions import IsAuthenticated
|
||
|
|
from rest_framework.response import Response
|
||
|
|
|
||
|
|
from apps.common.serializers import AiAnnotateTaskSerializer
|
||
|
|
from apps.core.models import (
|
||
|
|
AiAnnotateTask,
|
||
|
|
AiAnnotateTaskSample,
|
||
|
|
AiDatasetSample,
|
||
|
|
AIAnnotateTeam,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _get_labels_json_from_task(task: AiAnnotateTask) -> str:
|
||
|
|
if not task.labels:
|
||
|
|
return ""
|
||
|
|
try:
|
||
|
|
items = json.loads(task.labels)
|
||
|
|
except Exception:
|
||
|
|
return ""
|
||
|
|
if not isinstance(items, list):
|
||
|
|
return ""
|
||
|
|
return task.labels
|
||
|
|
|
||
|
|
|
||
|
|
def _build_text_prompt_from_labels_json(labels_json: str) -> str:
|
||
|
|
if not labels_json:
|
||
|
|
return ""
|
||
|
|
try:
|
||
|
|
items = json.loads(labels_json)
|
||
|
|
except Exception:
|
||
|
|
return ""
|
||
|
|
if isinstance(items, dict):
|
||
|
|
items = [items]
|
||
|
|
if not isinstance(items, list):
|
||
|
|
return ""
|
||
|
|
labels: List[str] = []
|
||
|
|
for it in items:
|
||
|
|
if isinstance(it, dict) and it.get("label"):
|
||
|
|
labels.append(str(it["label"]).strip())
|
||
|
|
labels = [x for x in labels if x]
|
||
|
|
return ".".join(labels) if labels else ""
|
||
|
|
|
||
|
|
|
||
|
|
def _flatten_bbox_region(region: Any) -> Optional[List[float]]:
|
||
|
|
try:
|
||
|
|
if (
|
||
|
|
isinstance(region, list)
|
||
|
|
and len(region) == 2
|
||
|
|
and isinstance(region[0], list)
|
||
|
|
and isinstance(region[1], list)
|
||
|
|
and len(region[0]) == 2
|
||
|
|
and len(region[1]) == 2
|
||
|
|
):
|
||
|
|
x1, y1 = float(region[0][0]), float(region[0][1])
|
||
|
|
x2, y2 = float(region[1][0]), float(region[1][1])
|
||
|
|
return [x1, y1, x2, y2]
|
||
|
|
except Exception:
|
||
|
|
return None
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _is_int_one(value: Any) -> bool:
|
||
|
|
try:
|
||
|
|
return int(str(value).strip()) == 1
|
||
|
|
except Exception:
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
def _run_auto_annotate_task(
|
||
|
|
task_id: str,
|
||
|
|
label_content: str,
|
||
|
|
text_prompt: str,
|
||
|
|
req_type: int,
|
||
|
|
score_threshold: float,
|
||
|
|
timeout_sec: float,
|
||
|
|
max_workers: int,
|
||
|
|
infer_url: str,
|
||
|
|
) -> None:
|
||
|
|
close_old_connections()
|
||
|
|
|
||
|
|
samples = list(AiAnnotateTaskSample.objects.filter(task_id=task_id, status="01"))
|
||
|
|
if not samples:
|
||
|
|
print(f"[auto-annotate] task_id={task_id} no samples with status=01")
|
||
|
|
close_old_connections()
|
||
|
|
return
|
||
|
|
|
||
|
|
health_url = infer_url.rstrip("/infer").rstrip("/") + "/health"
|
||
|
|
try:
|
||
|
|
resp = requests.get(health_url, timeout=5)
|
||
|
|
if resp.status_code != 200:
|
||
|
|
print(f"[auto-annotate] task_id={task_id} health check failed: {resp.status_code}")
|
||
|
|
close_old_connections()
|
||
|
|
return
|
||
|
|
print(f"[auto-annotate] task_id={task_id} health check passed")
|
||
|
|
except Exception as e:
|
||
|
|
print(f"[auto-annotate] task_id={task_id} health check error: {e}")
|
||
|
|
close_old_connections()
|
||
|
|
return
|
||
|
|
|
||
|
|
try:
|
||
|
|
from apps.common.minio_client import minio_storage
|
||
|
|
except Exception:
|
||
|
|
from common.minio_client import minio_storage
|
||
|
|
|
||
|
|
def call_infer(sample: AiAnnotateTaskSample) -> tuple[str, Optional[str], Optional[str]]:
|
||
|
|
sample_path = sample.sample_path or ""
|
||
|
|
sample_name = sample.sample_name or ""
|
||
|
|
|
||
|
|
if not sample_path or not sample_name:
|
||
|
|
return sample.id, None, "missing sample_path or sample_name"
|
||
|
|
|
||
|
|
object_name = os.path.join(sample_path.lstrip("/"), sample_name).replace("\\", "/")
|
||
|
|
|
||
|
|
try:
|
||
|
|
response = minio_storage.get_object("dataset", object_name)
|
||
|
|
image_data = response.read()
|
||
|
|
response.release_conn()
|
||
|
|
except Exception as e:
|
||
|
|
return sample.id, None, f"minio download error: {e}"
|
||
|
|
|
||
|
|
try:
|
||
|
|
resp = requests.post(
|
||
|
|
infer_url,
|
||
|
|
files={
|
||
|
|
"image_file": (sample_name, image_data, "image/jpeg"),
|
||
|
|
},
|
||
|
|
data={
|
||
|
|
"label_content": label_content,
|
||
|
|
"text_prompt": text_prompt,
|
||
|
|
"type": req_type,
|
||
|
|
"score_threshold": score_threshold,
|
||
|
|
},
|
||
|
|
timeout=timeout_sec,
|
||
|
|
)
|
||
|
|
except Exception as e:
|
||
|
|
return sample.id, None, str(e)
|
||
|
|
if resp.status_code != 200:
|
||
|
|
return sample.id, None, f"http {resp.status_code}: {resp.text[:300]}"
|
||
|
|
try:
|
||
|
|
result = resp.json()
|
||
|
|
except Exception:
|
||
|
|
return sample.id, None, f"invalid json: {resp.text[:300]}"
|
||
|
|
if not isinstance(result, list):
|
||
|
|
return sample.id, None, "invalid response: not list"
|
||
|
|
|
||
|
|
normalized: List[Dict[str, Any]] = []
|
||
|
|
for item in result:
|
||
|
|
if not isinstance(item, dict):
|
||
|
|
continue
|
||
|
|
if "label" not in item or "region" not in item:
|
||
|
|
continue
|
||
|
|
region = item.get("region")
|
||
|
|
if req_type == 0:
|
||
|
|
flat = _flatten_bbox_region(region)
|
||
|
|
if flat is None:
|
||
|
|
continue
|
||
|
|
out_region: Any = flat
|
||
|
|
else:
|
||
|
|
out_region = region
|
||
|
|
normalized.append(
|
||
|
|
{
|
||
|
|
"label": item.get("label"),
|
||
|
|
"name": item.get("name") or item.get("label"),
|
||
|
|
"color": item.get("color") or "red",
|
||
|
|
"region": out_region,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
return sample.id, json.dumps(normalized, ensure_ascii=False), None
|
||
|
|
|
||
|
|
start_time = time.time()
|
||
|
|
results: List[tuple[str, Optional[str], Optional[str]]] = []
|
||
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||
|
|
futures = [executor.submit(call_infer, s) for s in samples]
|
||
|
|
for fut in concurrent.futures.as_completed(futures):
|
||
|
|
results.append(fut.result())
|
||
|
|
|
||
|
|
ok_updates: List[AiAnnotateTaskSample] = []
|
||
|
|
errors: List[Dict[str, Any]] = []
|
||
|
|
now = datetime.now()
|
||
|
|
for sample_id, content, err in results:
|
||
|
|
if err or not content:
|
||
|
|
errors.append({"sample_id": sample_id, "error": err or "empty result"})
|
||
|
|
continue
|
||
|
|
ok_updates.append(
|
||
|
|
AiAnnotateTaskSample(
|
||
|
|
id=sample_id,
|
||
|
|
annotation_content=content,
|
||
|
|
status="02",
|
||
|
|
annotation_time=now,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
if ok_updates:
|
||
|
|
AiAnnotateTaskSample.objects.bulk_update(ok_updates, fields=["annotation_content", "status", "annotation_time"])
|
||
|
|
print(
|
||
|
|
f"[auto-annotate] task_id={task_id} success={len(ok_updates)} failed={len(errors)} "
|
||
|
|
f"elapsed={round(time.time() - start_time, 3)}s"
|
||
|
|
)
|
||
|
|
if errors:
|
||
|
|
print(f"[auto-annotate] task_id={task_id} errors={json.dumps(errors, ensure_ascii=False)}")
|
||
|
|
close_old_connections()
|
||
|
|
|
||
|
|
|
||
|
|
@api_view(["POST"])
|
||
|
|
def create_annotate_task(request):
|
||
|
|
data = request.data
|
||
|
|
data["id"] = str(uuid.uuid4())
|
||
|
|
data["status"] = "01"
|
||
|
|
data["creator"] = request.user.username if request.user.username else "creator"
|
||
|
|
data["create_time"] = datetime.now()
|
||
|
|
if _is_int_one(data.get("task_type")):
|
||
|
|
if _is_int_one(data.get("region_type")):
|
||
|
|
data["annotate_type"] = "03"
|
||
|
|
else:
|
||
|
|
data["annotate_type"] = "01"
|
||
|
|
serializer = AiAnnotateTaskSerializer(data=data, context={"request": request})
|
||
|
|
|
||
|
|
if serializer.is_valid():
|
||
|
|
serializer.save()
|
||
|
|
return Response({"message": "任务创建成功"})
|
||
|
|
|
||
|
|
return Response(
|
||
|
|
{"error": "无效的数据", "details": serializer.errors},
|
||
|
|
status=status.HTTP_400_BAD_REQUEST,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@api_view(["POST"])
|
||
|
|
@permission_classes([IsAuthenticated])
|
||
|
|
def start_auto_annotate_task(request):
|
||
|
|
data = request.data
|
||
|
|
task_id = data.get("task_id") or data.get("id")
|
||
|
|
if not task_id:
|
||
|
|
return Response({"error": "缺少 task_id"})
|
||
|
|
|
||
|
|
task = get_object_or_404(AiAnnotateTask, id=task_id)
|
||
|
|
user = request.user.username or "creator"
|
||
|
|
|
||
|
|
if not (task.task_team or task.task_leader):
|
||
|
|
return Response({"error": "未分配人员请进行分配"})
|
||
|
|
|
||
|
|
AiAnnotateTask.objects.filter(id=task_id).update(status="02", task_team=user, task_leader=user)
|
||
|
|
|
||
|
|
label_content = _get_labels_json_from_task(task)
|
||
|
|
if not label_content:
|
||
|
|
return Response({"error": "任务未配置 labels 或 labels 格式非法"})
|
||
|
|
|
||
|
|
text_prompt = str(data.get("text_prompt") or data.get("prompt") or "").strip()
|
||
|
|
if not text_prompt:
|
||
|
|
text_prompt = _build_text_prompt_from_labels_json(label_content)
|
||
|
|
if not text_prompt:
|
||
|
|
return Response({"error": "缺少 text_prompt"}, status=status.HTTP_400_BAD_REQUEST)
|
||
|
|
|
||
|
|
try:
|
||
|
|
req_type = int(data.get("type", int(task.region_type or "0")))
|
||
|
|
except Exception:
|
||
|
|
req_type = int(task.region_type or "0")
|
||
|
|
|
||
|
|
sample_count = AiAnnotateTaskSample.objects.filter(task_id=task_id).count()
|
||
|
|
if not sample_count:
|
||
|
|
return Response({"message": "当前任务没有样本", "success": 0, "failed": 0, "errors": []})
|
||
|
|
|
||
|
|
infer_url = getattr(settings, "ANNOTATION_SERVICE_INFER_URL", None) or ""
|
||
|
|
if not infer_url:
|
||
|
|
return Response({"error": "未配置 ANNOTATION_SERVICE_INFER_URL"}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
|
||
|
|
|
||
|
|
score_threshold = float(getattr(settings, "AUTO_ANNOTATE_SCORE_THRESHOLD", 0.4))
|
||
|
|
default_workers = int(getattr(settings, "AUTO_ANNOTATE_MAX_WORKERS", 2))
|
||
|
|
max_workers = int(data.get("max_workers") or default_workers)
|
||
|
|
max_workers = max(1, min(32, max_workers))
|
||
|
|
timeout_sec = float(data.get("timeout_sec") or 60)
|
||
|
|
t = threading.Thread(
|
||
|
|
target=_run_auto_annotate_task,
|
||
|
|
kwargs={
|
||
|
|
"task_id": str(task_id),
|
||
|
|
"label_content": str(label_content),
|
||
|
|
"text_prompt": str(text_prompt),
|
||
|
|
"req_type": int(req_type),
|
||
|
|
"score_threshold": float(score_threshold),
|
||
|
|
"timeout_sec": float(timeout_sec),
|
||
|
|
"max_workers": int(max_workers),
|
||
|
|
"infer_url": str(infer_url),
|
||
|
|
},
|
||
|
|
daemon=True,
|
||
|
|
)
|
||
|
|
t.start()
|
||
|
|
return Response(
|
||
|
|
{
|
||
|
|
"message": "自动标注任务已后台启动",
|
||
|
|
"task_id": task_id,
|
||
|
|
"type": req_type,
|
||
|
|
"sample_count": sample_count,
|
||
|
|
"max_workers": max_workers,
|
||
|
|
"text_prompt": text_prompt,
|
||
|
|
"thread_name": t.name,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@api_view(["GET"])
|
||
|
|
@permission_classes([IsAuthenticated])
|
||
|
|
def read_annotate_task(request):
|
||
|
|
task_id = request.GET.get("id", None)
|
||
|
|
task = get_object_or_404(AiAnnotateTask, id=task_id)
|
||
|
|
serializer = AiAnnotateTaskSerializer(task)
|
||
|
|
STATUSES = ["01", "02", "03", "04", "05", "06"]
|
||
|
|
sapmle_dict = {}
|
||
|
|
for status_code in STATUSES:
|
||
|
|
sapmle_dict[f"count_{status_code}"] = Sum(
|
||
|
|
Case(When(status=status_code, then=1), default=0, output_field=IntegerField())
|
||
|
|
)
|
||
|
|
sapmle_dict["total"] = Count("id")
|
||
|
|
task_status_counts = (
|
||
|
|
AiAnnotateTaskSample.objects.values("task_id")
|
||
|
|
.annotate(**sapmle_dict)
|
||
|
|
.filter(task_id=task_id)
|
||
|
|
)
|
||
|
|
serializer_data = serializer.data
|
||
|
|
serializer_data["task_status_counts"] = task_status_counts
|
||
|
|
return Response(serializer_data)
|
||
|
|
|
||
|
|
|
||
|
|
@api_view(["GET"])
|
||
|
|
@permission_classes([IsAuthenticated])
|
||
|
|
def get_mytask_annotate_list(request):
|
||
|
|
executor = request.user.username
|
||
|
|
sql = """
|
||
|
|
SELECT b.executor,a.task_name,a.task_type,a.region_type,a.id,a.status,TO_CHAR(a.create_time, 'YYYY-MM-DD') AS create_time,
|
||
|
|
COUNT(*) FILTER (WHERE b.status = '01') AS COUNT01,
|
||
|
|
COUNT(*) FILTER (WHERE b.status = '02') AS COUNT02,
|
||
|
|
COUNT(*) FILTER (WHERE b.status = '03') AS COUNT03,
|
||
|
|
COUNT(*) FILTER (WHERE b.status = '04') AS COUNT04,
|
||
|
|
COUNT(*) FILTER (WHERE b.status = '05') AS COUNT05,
|
||
|
|
COUNT(*) FILTER (WHERE b.status = '06') AS COUNT06,
|
||
|
|
COUNT(*) AS tota
|
||
|
|
FROM ai_annotate_task_sample b,ai_annotate_task a
|
||
|
|
WHERE a.id=b.task_id AND b.executor=%s and b.status in ('01','02','03','04','05','06')
|
||
|
|
GROUP BY b.executor,a.task_name,a.id,a.status,TO_CHAR(a.create_time, 'YYYY-MM-DD')
|
||
|
|
ORDER BY a.task_name
|
||
|
|
"""
|
||
|
|
with connection.cursor() as cursor:
|
||
|
|
cursor.execute(sql, [executor])
|
||
|
|
result = cursor.fetchall()
|
||
|
|
keys = [
|
||
|
|
"executor",
|
||
|
|
"task_name",
|
||
|
|
"task_type",
|
||
|
|
"region_type",
|
||
|
|
"id",
|
||
|
|
"status",
|
||
|
|
"create_time",
|
||
|
|
"COUNT01",
|
||
|
|
"COUNT02",
|
||
|
|
"COUNT03",
|
||
|
|
"COUNT04",
|
||
|
|
"COUNT05",
|
||
|
|
"COUNT06",
|
||
|
|
"total",
|
||
|
|
]
|
||
|
|
data = [dict(zip(keys, row)) for row in result]
|
||
|
|
return Response(data)
|
||
|
|
|
||
|
|
|
||
|
|
@api_view(["GET"])
|
||
|
|
@permission_classes([IsAuthenticated])
|
||
|
|
def get_mytask_audit_list(request):
|
||
|
|
executor = request.user.username
|
||
|
|
sql = """
|
||
|
|
SELECT
|
||
|
|
b.executor,
|
||
|
|
a.task_name,
|
||
|
|
a.id,
|
||
|
|
a.status,
|
||
|
|
TO_CHAR(a.create_time, 'YYYY-MM-DD') AS create_time,
|
||
|
|
COUNT(*) FILTER (WHERE b.status = '01') AS COUNT01,
|
||
|
|
COUNT(*) FILTER (WHERE b.status = '02') AS COUNT02,
|
||
|
|
COUNT(*) FILTER (WHERE b.status = '03') AS COUNT03,
|
||
|
|
COUNT(*) FILTER (WHERE b.status = '04') AS COUNT04,
|
||
|
|
COUNT(*) FILTER (WHERE b.status = '05') AS COUNT05,
|
||
|
|
COUNT(*) FILTER (WHERE b.status = '06') AS COUNT06,
|
||
|
|
COUNT(*) AS total
|
||
|
|
FROM
|
||
|
|
ai_annotate_task_sample b,
|
||
|
|
ai_annotate_task a
|
||
|
|
WHERE
|
||
|
|
a.id=b.task_id AND a.task_leader=%s and b.status in ('01','02','03','04','05','06')
|
||
|
|
GROUP BY
|
||
|
|
b.executor, a.task_name, a.id, a.status, TO_CHAR(a.create_time, 'YYYY-MM-DD')
|
||
|
|
ORDER BY
|
||
|
|
a.task_name
|
||
|
|
"""
|
||
|
|
|
||
|
|
with connection.cursor() as cursor:
|
||
|
|
cursor.execute(sql, [executor])
|
||
|
|
result = cursor.fetchall()
|
||
|
|
|
||
|
|
keys = [
|
||
|
|
"executor",
|
||
|
|
"task_name",
|
||
|
|
"id",
|
||
|
|
"status",
|
||
|
|
"create_time",
|
||
|
|
"COUNT01",
|
||
|
|
"COUNT02",
|
||
|
|
"COUNT03",
|
||
|
|
"COUNT04",
|
||
|
|
"COUNT05",
|
||
|
|
"COUNT06",
|
||
|
|
"total",
|
||
|
|
]
|
||
|
|
data = [dict(zip(keys, row)) for row in result]
|
||
|
|
return Response(data)
|
||
|
|
|
||
|
|
|
||
|
|
@api_view(["PUT"])
|
||
|
|
@permission_classes([IsAuthenticated])
|
||
|
|
def update_annotate_task(request):
|
||
|
|
data = request.data
|
||
|
|
task = get_object_or_404(AiAnnotateTask, id=data.get("id"))
|
||
|
|
task_type = data.get("task_type", task.task_type)
|
||
|
|
region_type = data.get("region_type", task.region_type)
|
||
|
|
if _is_int_one(task_type):
|
||
|
|
if _is_int_one(region_type):
|
||
|
|
data["annotate_type"] = "03"
|
||
|
|
else:
|
||
|
|
data["annotate_type"] = "01"
|
||
|
|
serializer = AiAnnotateTaskSerializer(instance=task, data=data, partial=True)
|
||
|
|
|
||
|
|
if serializer.is_valid():
|
||
|
|
serializer.save()
|
||
|
|
return Response({"message": "任务更新成功"})
|
||
|
|
|
||
|
|
return Response(
|
||
|
|
{"error": "无效的数据", "details": serializer.errors},
|
||
|
|
status=status.HTTP_400_BAD_REQUEST,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@api_view(["DELETE"])
|
||
|
|
@permission_classes([IsAuthenticated])
|
||
|
|
def delete_annotate_task(request):
|
||
|
|
task_id = request.GET.get("id", None)
|
||
|
|
task = get_object_or_404(AiAnnotateTask, id=task_id)
|
||
|
|
task.delete()
|
||
|
|
return Response({"message": "任务删除成功"})
|
||
|
|
|
||
|
|
|
||
|
|
@api_view(["GET"])
|
||
|
|
@permission_classes([IsAuthenticated])
|
||
|
|
def search_annotate_tasks(request):
|
||
|
|
algorithm_id = request.GET.get("algorithm_id", None)
|
||
|
|
status_code = request.GET.get("status", None)
|
||
|
|
custom1 = request.GET.get("custom1", None)
|
||
|
|
page_size = request.GET.get("page_size", 10)
|
||
|
|
filter_kwargs = {}
|
||
|
|
if algorithm_id:
|
||
|
|
filter_kwargs["algorithm_id"] = algorithm_id
|
||
|
|
if status_code:
|
||
|
|
filter_kwargs["status"] = status_code
|
||
|
|
if custom1:
|
||
|
|
filter_kwargs["custom1__contains"] = custom1
|
||
|
|
|
||
|
|
tasks = AiAnnotateTask.objects.filter(**filter_kwargs).order_by("task_name")
|
||
|
|
paginator = PageNumberPagination()
|
||
|
|
paginator.page_size = page_size
|
||
|
|
result_page = paginator.paginate_queryset(tasks, request)
|
||
|
|
serializer = AiAnnotateTaskSerializer(result_page, many=True)
|
||
|
|
|
||
|
|
response_data = serializer.data
|
||
|
|
for task_data in response_data:
|
||
|
|
task_id = task_data["id"]
|
||
|
|
sample_count = AiAnnotateTaskSample.objects.filter(task_id=task_id).count()
|
||
|
|
audited_count = AiAnnotateTaskSample.objects.filter(
|
||
|
|
task_id=task_id, status__in=["04", "06"]
|
||
|
|
).count()
|
||
|
|
|
||
|
|
task_data["sample_count"] = sample_count
|
||
|
|
task_data["audited_count"] = audited_count
|
||
|
|
task_data["complete_percent"] = (
|
||
|
|
round(audited_count / sample_count, 2) if sample_count != 0 else 0
|
||
|
|
)
|
||
|
|
|
||
|
|
return paginator.get_paginated_response(response_data)
|