import os import uuid from datetime import datetime import xml.etree.ElementTree as ET from django.db import transaction from django.db.models import Case, Count, IntegerField, Q, 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 rest_framework.utils import json from apps.common.serializers import AiAnnotateTaskSampleSerializer from apps.common.utils.annotation_store import ( get_dataset_sample_annotation_content, split_dataset_sample_annotation_content, ) from apps.core.models import ( AiAnnotateTask, AiAnnotateTaskSample, AiDatasetSample, ) from apps.datasets.dataset_stats import refresh_dataset_sample_counters from apps.datasets.storage_service import download_sidecar_text, upload_sidecar_bytes def _safe_split_wh(custom1): try: w_str, h_str = str(custom1).lower().split("x") return int(w_str), int(h_str) except Exception: return None, None def _extract_sample_index(saved_filename): stem = os.path.splitext(os.path.basename(str(saved_filename or "")))[0] tail = stem.rsplit("-", 1)[-1] if tail.isdigit(): return int(tail) return None def _get_dataset_sample_summary_status(dataset_sample, task_sample_statuses): if dataset_sample.rectangle or dataset_sample.polygon or "06" in task_sample_statuses: return "02" if task_sample_statuses: return "01" return "00" def _refresh_dataset_sample_status(sample_ids): if not sample_ids: return sample_ids = list({sample_id for sample_id in sample_ids if sample_id}) samples = list( AiDatasetSample.objects.filter(id__in=sample_ids).only("id", "status", "rectangle", "polygon") ) if not samples: return task_sample_status_map = {} for sample_id, status_code in AiAnnotateTaskSample.objects.filter(sample_id__in=sample_ids).values_list( "sample_id", "status" ): task_sample_status_map.setdefault(sample_id, set()).add(status_code) objects_to_update = [] for sample in samples: new_status = _get_dataset_sample_summary_status( sample, task_sample_status_map.get(sample.id, set()) ) if sample.status != new_status: objects_to_update.append(AiDatasetSample(id=sample.id, status=new_status)) if objects_to_update: AiDatasetSample.objects.bulk_update(objects_to_update, fields=["status"]) def _normalize_polygon_points(region): if not isinstance(region, list): return [] pts = [] for p in region: if not isinstance(p, (list, tuple)) or len(p) != 2: continue try: x = float(p[0]) y = float(p[1]) except Exception: continue pts.append([x, y]) return pts if len(pts) >= 3 else [] def _shape_key(label, points): try: norm = [(round(float(x), 2), round(float(y), 2)) for x, y in points] except Exception: return None return str(label), json.dumps(norm, separators=(",", ":"), ensure_ascii=False) def _is_polygon_items(items): if not isinstance(items, list): return False for it in items: if not isinstance(it, dict): continue region = it.get("region") if isinstance(region, list) and region and isinstance(region[0], (list, tuple)): return True return False def _bbox_key(item): try: label = str(item.get("label")) region = item.get("region") if not isinstance(region, list) or len(region) != 4: return None vals = [round(float(x), 2) for x in region] return label, json.dumps(vals, separators=(",", ":"), ensure_ascii=False) except Exception: return None def _shapes_from_items(items): if not isinstance(items, list): return [] shapes = [] for it in items: if not isinstance(it, dict): continue label = it.get("label") if not label: continue points = _normalize_polygon_points(it.get("region")) if not points: continue shapes.append( { "label": str(label), "points": points, "group_id": None, "shape_type": "polygon", "flags": {}, "mask": None, } ) return shapes def _shapes_from_labelme(labelme_obj): if not isinstance(labelme_obj, dict): return [] shapes_obj = labelme_obj.get("shapes") if not isinstance(shapes_obj, list): return [] shapes = [] for sh in shapes_obj: if not isinstance(sh, dict): continue if sh.get("shape_type") != "polygon": continue label = sh.get("label") points = _normalize_polygon_points(sh.get("points")) if not label or not points: continue shapes.append( { "label": str(label), "points": points, "group_id": None, "shape_type": "polygon", "flags": {}, "mask": None, } ) return shapes def _load_polygon_shapes_from_dataset_sample(dataset_sample): content = get_dataset_sample_annotation_content(dataset_sample) if not content: return [] try: parsed = json.loads(content) except Exception: return [] if isinstance(parsed, dict) and parsed.get("type") == "polygon": try: json_text = download_sidecar_text(dataset_sample.saved_path, dataset_sample.saved_filename, ".json") if not json_text: return [] labelme_obj = json.loads(json_text) except Exception: return [] return _shapes_from_labelme(labelme_obj) if isinstance(parsed, list): return _shapes_from_items(parsed) return [] @api_view(["POST"]) @permission_classes([IsAuthenticated]) def set_person_task(request): data = request.data task_id = data["task_id"] task = AiAnnotateTask.objects.get(id=task_id) dataset_id = task.dataset_id status_code = data.get("status") task_range = task.task_range query_condition = Q(dataset_id=dataset_id) if status_code: query_condition &= Q(status=status_code) all_samples = list(AiDatasetSample.objects.filter(query_condition).order_by("saved_filename")) existing_sample_ids = set( AiAnnotateTaskSample.objects.filter(task_id=task_id).values_list("sample_id", flat=True) ) if existing_sample_ids: all_samples = [sample for sample in all_samples if sample.id not in existing_sample_ids] if task_range: start, end = map(int, task_range.split("-")) all_samples = [ sample for sample in all_samples if (sample_index := _extract_sample_index(sample.saved_filename)) is not None and start <= sample_index <= end ] all_samples_length = len(all_samples) rule_type = data["rule_type"] task_rules = json.loads(data["task_rules"]) total_percent = 0 last_person = "" for key, value in task_rules.items(): person = key if rule_type == "percent": percent = value total_percent += percent last_person = person num_samples_to_assign = int(all_samples_length * percent / 100) create_person_samples(task_id, person, all_samples[:num_samples_to_assign]) all_samples = all_samples[num_samples_to_assign:] elif rule_type == "number": number = value create_person_samples(task_id, person, all_samples[:number]) all_samples = all_samples[number:] elif rule_type == "range": start, end = map(int, value.split("-")) matching_samples = [ sample for sample in all_samples if (sample_index := _extract_sample_index(sample.saved_filename)) is not None and start <= sample_index <= end ] create_person_samples(task_id, person, matching_samples) all_samples = list(set(all_samples) - set(matching_samples)) if rule_type == "percent": remaining_samples = len(all_samples) if total_percent == 100 and remaining_samples > 0: create_person_samples(task_id, last_person, all_samples) return Response({"message": "人员标注任务分配成功"}) @transaction.atomic def create_person_samples(task_id, person, assigned_samples): if not assigned_samples: return objects_to_save = [ AiAnnotateTaskSample( id=str(uuid.uuid4()), task_id=task_id, executor=person, sample_id=sample.id, sample_path=sample.saved_path, sample_name=sample.saved_filename, status="01", custom1=sample.custom1, ) for sample in assigned_samples ] AiAnnotateTaskSample.objects.bulk_create(objects_to_save) _refresh_dataset_sample_status([sample.id for sample in assigned_samples]) @api_view(["POST"]) @permission_classes([IsAuthenticated]) def annotate_task_sample(request): data = request.data sample_id = data.get("id") hard = data.get("hard") task_sample = get_object_or_404(AiAnnotateTaskSample, id=sample_id) AiAnnotateTaskSample.objects.filter(id=sample_id).update( annotation_content=data.get("annotation_content"), status="02", hard=hard, annotation_time=datetime.now(), ) _refresh_dataset_sample_status([task_sample.sample_id]) return Response({"message": "ok"}) @api_view(["POST"]) @permission_classes([IsAuthenticated]) def submit_task_samples(request): data = request.data task_id = data.get("task_id") person = data.get("person") ids = [id for id in data.get("ids", "").split(",") if id] queryset = AiAnnotateTaskSample.objects.filter(task_id=task_id, executor=person, status="02") if ids: queryset = queryset.filter(id__in=ids) sample_ids = list(queryset.values_list("sample_id", flat=True)) queryset.update(status="03") _refresh_dataset_sample_status(sample_ids) return Response({"message": "ok"}) @api_view(["GET"]) @permission_classes([IsAuthenticated]) def read_task_sample(request): sample_id = request.GET.get("id", None) sample = get_object_or_404(AiAnnotateTaskSample, id=sample_id) serializer = AiAnnotateTaskSampleSerializer(sample) return Response(serializer.data) @api_view(["POST"]) @permission_classes([IsAuthenticated]) def audit_task_samples(request): data = request.data task_id = data.get("task_id") person = request.user.username feedback = data.get("feedback") status_code = data.get("status") ids = [id for id in data.get("ids", "").split(",") if id] queryset = AiAnnotateTaskSample.objects.filter(task_id=task_id, status="03") if ids: queryset = queryset.filter(id__in=ids) if status_code == "04": sample_ids = list(queryset.values_list("sample_id", flat=True)) queryset.update(status="04", auditor=person) _refresh_dataset_sample_status(sample_ids) elif status_code == "05": sample_ids = list(queryset.values_list("sample_id", flat=True)) queryset.update(status="05", auditor=person, feedback=feedback) _refresh_dataset_sample_status(sample_ids) return Response({"message": "ok"}) @api_view(["POST"]) @permission_classes([IsAuthenticated]) @transaction.atomic def merge_sample_content(request): try: data = request.data task_id = data.get("task_id") if not task_id: return Response({"error": "缺少 task_id"}, status=status.HTTP_400_BAD_REQUEST) datasettask = AiAnnotateTask.objects.get(id=task_id) executor = data.get("executor", None) filter_kwargs = {"task_id": task_id} if executor: filter_kwargs["executor"] = executor filter_kwargs["status"] = "04" queryset = AiAnnotateTaskSample.objects.filter(**filter_kwargs) samples = list(queryset) sample_ids = [sample.sample_id for sample in samples] with transaction.atomic(): for sample in samples: if sample.annotation_content: person_items = json.loads(sample.annotation_content) dataset_sample = AiDatasetSample.objects.get(id=sample.sample_id) if datasettask.annotate_type == "03" or _is_polygon_items(person_items): annotation_fields = split_dataset_sample_annotation_content(sample.annotation_content) AiDatasetSample.objects.filter(id=sample.sample_id).update( rectangle=annotation_fields["rectangle"], polygon=annotation_fields["polygon"], ) else: dataset_annotation = get_dataset_sample_annotation_content(dataset_sample) dataset_items = json.loads(dataset_annotation) if dataset_annotation else [] merged_data = person_items + dataset_items unique_annotations = {} for item in merged_data: key = _bbox_key(item) if not key: continue if key not in unique_annotations: unique_annotations[key] = item final_merged_data = list(unique_annotations.values()) imgwidth, imgheight = _safe_split_wh(sample.custom1) xml_bytes = _build_voc_xml_bytes( final_merged_data, sample.sample_path, sample.sample_name, imgwidth, imgheight, ) upload_sidecar_bytes( sample.sample_path, sample.sample_name, ".xml", xml_bytes, "application/xml", ) merged_json = json.dumps(final_merged_data, ensure_ascii=False) annotation_fields = split_dataset_sample_annotation_content(merged_json) AiDatasetSample.objects.filter(id=sample.sample_id).update( rectangle=annotation_fields["rectangle"], polygon=annotation_fields["polygon"], ) queryset.update(status="06") _refresh_dataset_sample_status(sample_ids) refresh_dataset_sample_counters(datasettask.dataset_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() if sample_count == audited_count: AiAnnotateTask.objects.filter(id=task_id).update(status="03") return Response({"message": "任务样本成功合并到数据集中"}) except Exception as e: return Response( {"error": "合并入库失败", "details": str(e)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR, ) def _build_voc_xml_bytes(data, sample_path, filename, width, height): normalized_path = str(sample_path or "").strip("/").replace("\\", "/") image_path = f"{normalized_path}/{filename}" if normalized_path else filename annotation = ET.Element("annotation") folder = ET.SubElement(annotation, "folder") folder.text = normalized_path filename_elem = ET.SubElement(annotation, "filename") filename_elem.text = filename path = ET.SubElement(annotation, "path") path.text = image_path source = ET.SubElement(annotation, "source") database = ET.SubElement(source, "database") database.text = "Unknown" size = ET.SubElement(annotation, "size") width_elem = ET.SubElement(size, "width") width_elem.text = str(int(width) if width is not None else 0) height_elem = ET.SubElement(size, "height") height_elem.text = str(int(height) if height is not None else 0) depth = ET.SubElement(size, "depth") depth.text = "3" segmented = ET.SubElement(annotation, "segmented") segmented.text = "0" for obj in data: object_elem = ET.SubElement(annotation, "object") name = ET.SubElement(object_elem, "name") name.text = obj["label"] pose = ET.SubElement(object_elem, "pose") pose.text = "Unspecified" truncated = ET.SubElement(object_elem, "truncated") truncated.text = "0" difficult = ET.SubElement(object_elem, "difficult") difficult.text = "0" bndbox = ET.SubElement(object_elem, "bndbox") xmin, ymin, xmax, ymax = map(float, obj["region"]) xmin_elem = ET.SubElement(bndbox, "xmin") xmin_elem.text = str(int(xmin)) ymin_elem = ET.SubElement(bndbox, "ymin") ymin_elem.text = str(int(ymin)) xmax_elem = ET.SubElement(bndbox, "xmax") xmax_elem.text = str(int(xmax)) ymax_elem = ET.SubElement(bndbox, "ymax") ymax_elem.text = str(int(ymax)) return ET.tostring(annotation, encoding="utf-8", xml_declaration=True) @api_view(["DELETE"]) @permission_classes([IsAuthenticated]) @transaction.atomic def delete_task_samples(request): data = request.data task_id = data.get("task_id") person = data.get("person") ids = [id for id in data.get("ids", "").split(",") if id] try: queryset = AiAnnotateTaskSample.objects.filter(task_id=task_id, executor=person) if ids: queryset = queryset.filter(id__in=ids) samples = list(queryset) except AiAnnotateTaskSample.DoesNotExist: return Response({"message": "没有选择对应的样本数据！"}, status=status.HTTP_404_NOT_FOUND) with transaction.atomic(): sample_ids = [sample.sample_id for sample in samples] queryset.delete() _refresh_dataset_sample_status(sample_ids) return Response({"message": "任务样本删除成功"}) @api_view(["GET"]) @permission_classes([IsAuthenticated]) def search_task_samples(request): task_id = request.GET.get("task_id", None) executor = request.GET.get("executor", None) status_code = request.GET.get("status", "") status_in = status_code.split(",") if status_code else [] sample_name = request.GET.get("sample_name", None) is_hard = request.GET.get("hard", None) page_size = request.GET.get("page_size", 10) filter_kwargs = {} if task_id: filter_kwargs["task_id"] = task_id if executor: filter_kwargs["executor"] = executor if status_code: filter_kwargs["status__in"] = status_in if sample_name: filter_kwargs["sample_name__contains"] = sample_name if is_hard: filter_kwargs["hard"] = is_hard samples = AiAnnotateTaskSample.objects.filter(**filter_kwargs).order_by("sample_name") paginator = PageNumberPagination() paginator.page_size = page_size result_page = paginator.paginate_queryset(samples, request) formatted_samples = [] for sample in result_page: sample.annotation_time = ( sample.annotation_time.strftime("%Y-%m-%d %H:%M:%S") if sample.annotation_time else None ) formatted_samples.append(sample) serializer = AiAnnotateTaskSampleSerializer(formatted_samples, many=True) return paginator.get_paginated_response(serializer.data) @api_view(["GET"]) @permission_classes([IsAuthenticated]) def query_person_tasks(request): task_id = request.GET.get("task_id", None) STATUSES = ["01", "02", "03", "04", "05", "06"] annotate_dict = {} for status_code in STATUSES: annotate_dict[f"count_{status_code}"] = Sum( Case(When(status=status_code, then=1), default=0, output_field=IntegerField()) ) annotate_dict["total"] = Count("id") executor_status_counts = ( AiAnnotateTaskSample.objects.values("executor") .annotate(**annotate_dict) .filter(task_id=task_id) .order_by("executor") ) return Response(executor_status_counts)