import json import io import logging import os import tarfile import threading import time import uuid import zipfile from datetime import datetime import xml.etree.ElementTree as ET from typing import List from django.conf import settings from django.db import close_old_connections, connection, transaction from django.db.models import Case, Count, IntegerField, Max, Q, Sum, When from django.shortcuts import get_object_or_404 from rest_framework import status from rest_framework.decorators import api_view, parser_classes, permission_classes from rest_framework.pagination import PageNumberPagination from rest_framework.parsers import MultiPartParser from rest_framework.permissions import IsAuthenticated from rest_framework.response import Response from apps.common.serializers import AiDatasetSerializer, AiDatasetSampleSerializer from apps.core.models import AiDataset from apps.core.models import ( AiAlgorithmTrainRecords, AiDatasetSample, AiAnnotateTask, AiTempFile, ) from apps.common.minio_client import minio_storage from apps.common.utils.annotation_store import split_dataset_sample_annotation_content from apps.datasets.dataset_stats import refresh_dataset_sample_counters from apps.datasets.storage_service import ( VIDEO_EXTENSIONS, copy_minio_object_to_dataset_result, copy_minio_object_to_dataset, dataset_type_accepts_extension, extract_video_frames_to_dataset_result, extract_video_frames_to_dataset, import_zip_bytes_to_dataset_result, import_zip_bytes_to_dataset, save_bytes_to_dataset, upload_sidecar_bytes, ) IMAGE_EXTENSIONS = ("jpg", "jpeg", "png", "bmp", "gif") logger = logging.getLogger(__name__) def _build_filespace_url(rel_path: str) -> str: base = str(getattr(settings, "IMAGE_BASE_URL", "") or "").rstrip("/") rel = str(rel_path or "").lstrip("/").replace("\\", "/") return f"{base}/{rel}" if base else rel def _strip_nginx_image_prefix(raw: str) -> str: s = str(raw or "").strip().replace("\\", "/") if not s: return "" if "://" not in s: return s.lstrip("/") try: from urllib.parse import urlsplit parts = urlsplit(s) path = (parts.path or "").lstrip("/") except Exception: idx = s.find("://") path = s[idx + 3 :] slash = path.find("/") path = path[slash + 1 :] if slash >= 0 else "" path = path.lstrip("/") if path.startswith("image/"): path = path[len("image/") :] return path.lstrip("/") def _videos_base_rel(videos_dir_rel: str) -> str: s = str(videos_dir_rel or "").strip().strip("/").replace("\\", "/") return f"{s}/" if s else "" def _safe_load_json_list(raw: str) -> List[dict]: if not raw: return [] s = str(raw).strip() if not s.startswith("["): return [] try: obj = json.loads(s) except Exception: return [] if isinstance(obj, list): return [x for x in obj if isinstance(x, dict)] return [] def _safe_load_json_value(raw, default): if raw in (None, ""): return default if isinstance(raw, (dict, list)): return raw try: return json.loads(raw) except Exception: return default def _find_dataset_sample_for_xml(dataset_id: str, image_name: str, xml_name: str): queryset = AiDatasetSample.objects.filter(dataset_id=dataset_id).order_by("-create_time") candidates = [] for name in (image_name, os.path.basename(image_name or ""), xml_name): s = str(name or "").strip() if s: candidates.append(os.path.basename(s)) for candidate in candidates: exact = queryset.filter(Q(original_filename=candidate) | Q(saved_filename=candidate)).first() if exact: return exact stems = {os.path.splitext(name)[0].lower() for name in candidates if name} if not stems: return None for sample in queryset: original_stem = os.path.splitext(str(sample.original_filename or ""))[0].lower() saved_stem = os.path.splitext(str(sample.saved_filename or ""))[0].lower() if original_stem in stems or saved_stem in stems: return sample return None def _parse_voc_xml_annotations(file_content: bytes, dataset_labels) -> tuple[str, list[dict]]: root = ET.fromstring(file_content) image_name = root.findtext("filename") or "" label_config = _safe_load_json_value(dataset_labels, []) label_map = {} for label in label_config if isinstance(label_config, list) else []: if not isinstance(label, dict): continue code = str(label.get("label") or "").strip() name = str(label.get("name") or code).strip() if code: label_map[code] = { "label": code, "name": name or code, "color": label.get("color") or "#409eff", } result_list = [] for obj in root.findall("object"): obj_name = str(obj.findtext("name") or "").strip() bndbox = obj.find("bndbox") if not obj_name or bndbox is None: continue xmin = bndbox.findtext("xmin") ymin = bndbox.findtext("ymin") xmax = bndbox.findtext("xmax") ymax = bndbox.findtext("ymax") if None in {xmin, ymin, xmax, ymax}: continue matched_label = label_map.get( obj_name, {"label": obj_name, "name": obj_name, "color": "#409eff"}, ) result_list.append( { "label": matched_label["label"], "name": matched_label["name"], "color": matched_label["color"], "region": [xmin, ymin, xmax, ymax], } ) return image_name, result_list def _load_temp_file_imported_dataset_ids(raw: str) -> list[str]: return [item for item in str(raw or "").split(",") if item] def _dump_temp_file_imported_dataset_ids(dataset_ids: list[str]) -> str: return ",".join(dict.fromkeys([item for item in dataset_ids if item])) def _compact_temp_import_summary(payload: dict, max_length: int = 1000) -> str: candidate = dict(payload or {}) for limit in (5, 3, 1, 0): shrunk = dict(candidate) if isinstance(shrunk.get("failures"), list): failures = shrunk["failures"][:limit] shrunk["failures"] = failures shrunk["failures_truncated"] = len(candidate.get("failures", [])) > len(failures) if isinstance(shrunk.get("duplicates"), list): duplicates = shrunk["duplicates"][:limit] shrunk["duplicates"] = duplicates shrunk["duplicates_truncated"] = len(candidate.get("duplicates", [])) > len(duplicates) text = json.dumps(shrunk, ensure_ascii=False, separators=(",", ":")) if len(text) <= max_length: return text return json.dumps({"status": payload.get("status"), "message": payload.get("message")}, ensure_ascii=False)[:max_length] def _normalize_temp_file_ids(data) -> list[str]: raw_values = [] if hasattr(data, "getlist"): raw_values.extend(data.getlist("temp_file_ids")) temp_file_ids = data.get("temp_file_ids") if isinstance(temp_file_ids, (list, tuple)): raw_values.extend(temp_file_ids) elif temp_file_ids not in (None, ""): raw_values.extend(str(temp_file_ids).split(",")) temp_file_id = data.get("temp_file_id") if temp_file_id not in (None, ""): raw_values.append(temp_file_id) normalized = [] seen = set() for item in raw_values: value = str(item or "").strip() if not value or value in seen: continue normalized.append(value) seen.add(value) return normalized def _import_single_temp_file_to_dataset(dataset: AiDataset, temp_file: AiTempFile, interval_seconds: int, max_frames: int) -> tuple[int, dict]: dataset_id = str(dataset.id) with transaction.atomic(): bucket_name = temp_file.bucket_name or "tempfile" object_name = temp_file.object_name if not object_name: return status.HTTP_400_BAD_REQUEST, {"error": "临时文件对象不存在"} info = minio_storage.get_object_info(bucket_name, object_name) filename = temp_file.original_name or os.path.basename(object_name) original_type = str(temp_file.original_type or "") imported_dataset_ids = _load_temp_file_imported_dataset_ids(temp_file.custom3) if dataset_id in imported_dataset_ids: latest_summary = _safe_load_json_value(temp_file.custom2, {}) return status.HTTP_409_CONFLICT, { "error": "该临时文件已导入过当前数据集，无需重复导入", "latest_import": latest_summary, } source_marker_prefix = f"temp:{temp_file.id}" result = {"imported": [], "skipped": [], "failed": []} if original_type == "04": zip_bytes = minio_storage.download_bytes(bucket_name, object_name) result = import_zip_bytes_to_dataset_result( dataset=dataset, zip_bytes=zip_bytes, source_marker_prefix=source_marker_prefix, skip_existing=True, ) elif dataset.dataset_type == "02" and original_type == "02": video_bytes = minio_storage.download_bytes(bucket_name, object_name) result = extract_video_frames_to_dataset_result( dataset=dataset, video_bytes=video_bytes, original_filename=filename, interval_seconds=interval_seconds, max_frames=max_frames, source_marker_prefix=source_marker_prefix, skip_existing=True, ) else: ext = os.path.splitext(filename)[1].lower() if not dataset_type_accepts_extension(dataset.dataset_type, ext): return status.HTTP_400_BAD_REQUEST, {"error": "临时文件类型与数据集类型不匹配"} result = copy_minio_object_to_dataset_result( dataset=dataset, source_bucket=bucket_name, source_object=object_name, original_filename=filename, size_bytes=info.get("size", 0), source_marker=f"{source_marker_prefix}:object", skip_existing=True, ) imported_samples = result["imported"] duplicate_items = result["skipped"] failed_items = result["failed"] imported_count = len(imported_samples) duplicate_count = len(duplicate_items) failed_count = len(failed_items) if imported_count > 0: refresh_dataset_sample_counters(dataset_id) summary_status = "success" if imported_count > 0 and failed_count > 0: summary_status = "partial" elif imported_count == 0 and duplicate_count > 0 and failed_count == 0: summary_status = "duplicate" elif imported_count == 0 and failed_count > 0: summary_status = "failed" summary_payload = { "status": summary_status, "dataset_id": dataset_id, "temp_file_id": temp_file.id, "bucket_name": bucket_name, "object_name": object_name, "filename": filename, "original_type": original_type, "etag": info.get("etag"), "imported_count": imported_count, "duplicate_count": duplicate_count, "failed_count": failed_count, "duplicates": duplicate_items, "failures": failed_items, } temp_file.custom1 = summary_status temp_file.custom2 = _compact_temp_import_summary(summary_payload) if imported_count > 0 and failed_count == 0: imported_dataset_ids.append(dataset_id) temp_file.custom3 = _dump_temp_file_imported_dataset_ids(imported_dataset_ids) if imported_count > 0: temp_file.status = "02" temp_file.save(update_fields=["status", "custom1", "custom2", "custom3"]) if failed_items: logger.warning("temp file import partial/failed: %s", summary_payload) else: logger.info("temp file import success: %s", summary_payload) if imported_count == 0 and duplicate_count > 0 and failed_count == 0: return status.HTTP_409_CONFLICT, { "message": "临时文件已导入过当前数据集，本次未新增样本", "imported_count": 0, "duplicate_count": duplicate_count, "failed_count": 0, "duplicates": duplicate_items, "failures": [], "sample_ids": [], } if imported_count == 0 and failed_count > 0: return status.HTTP_400_BAD_REQUEST, { "error": "未导入任何样本，请检查文件类型与数据集类型是否匹配", "imported_count": 0, "duplicate_count": duplicate_count, "failed_count": failed_count, "duplicates": duplicate_items, "failures": failed_items, "sample_ids": [], } message = "临时文件导入数据集成功" if failed_count > 0: message = f"临时文件导入部分成功，成功{imported_count}条，失败{failed_count}条" elif duplicate_count > 0: message = f"临时文件导入完成，新增{imported_count}条，跳过重复{duplicate_count}条" return status.HTTP_200_OK, { "message": message, "imported_count": imported_count, "duplicate_count": duplicate_count, "failed_count": failed_count, "duplicates": duplicate_items, "failures": failed_items, "sample_ids": [sample.id for sample in imported_samples], } def _load_video_manifest(raw: str) -> tuple[str, List[dict]]: if not raw: return "", [] s = str(raw).strip() if not s: return "", [] if s.startswith("{"): try: obj = json.loads(s) except Exception: return "", [] if not isinstance(obj, dict): return "", [] base = str(obj.get("b") or obj.get("base") or "") items = obj.get("v") or obj.get("videos") or [] if isinstance(items, list): vids = [x for x in items if isinstance(x, dict)] else: vids = [] return base, vids if s.startswith("["): try: obj = json.loads(s) except Exception: return "", [] if not isinstance(obj, list): return "", [] vids = [x for x in obj if isinstance(x, dict)] return "", vids return "", [] def _dump_video_manifest(base_url: str, videos: List[dict], max_len: int = 1000) -> str: b = str(base_url or "") keep = list(videos or []) while True: payload = {"b": b, "v": keep} s = json.dumps(payload, ensure_ascii=False, separators=(",", ":")) if len(s) <= max_len: return s if len(keep) <= 1: break keep = keep[-max(1, len(keep) // 2) :] payload = {"b": b, "v": keep} s = json.dumps(payload, ensure_ascii=False, separators=(",", ":")) if len(s) <= max_len: return s overhead = len(json.dumps({"b": "", "v": keep}, ensure_ascii=False, separators=(",", ":"))) allow_b = max(0, max_len - overhead) b2 = b[:allow_b] return json.dumps({"b": b2, "v": keep}, ensure_ascii=False, separators=(",", ":")) def _get_videos_dir_rel(dataset: AiDataset) -> str: if dataset.custom3: return dataset.custom3 base_dir = os.path.dirname((dataset.original_file_path or "").rstrip("/\\")) return os.path.join(base_dir, "video").replace("\\", "/") if base_dir else "" def _build_dataset_root(dataset_code: str) -> str: return str(dataset_code or "").strip().replace("\\", "/").strip("/") def _build_dataset_paths(dataset_code: str) -> dict: root = _build_dataset_root(dataset_code) return { "original_file_path": os.path.join(root, "file").replace("\\", "/"), "dataset_path": os.path.join(root, "image").replace("\\", "/"), "custom1": os.path.join(root, "image_mask").replace("\\", "/"), "custom3": os.path.join(root, "video").replace("\\", "/"), } def _generate_next_dataset_code() -> str: # Lock the dataset table so concurrent creates cannot generate the same code. with connection.cursor() as cursor: cursor.execute("LOCK TABLE ai_dataset IN EXCLUSIVE MODE") cursor.execute( """ SELECT COALESCE(MAX(CAST(dataset_code AS INTEGER)), 0) FROM ai_dataset WHERE dataset_code ~ '^[0-9]+$' """ ) next_index = int((cursor.fetchone() or [0])[0] or 0) + 1 return f"{next_index:06d}" def _ensure_video_storage(dataset: AiDataset) -> str: videos_dir_rel = _get_videos_dir_rel(dataset) if not videos_dir_rel: return "" if dataset.custom3 != videos_dir_rel: dataset.custom3 = videos_dir_rel dataset.save(update_fields=["custom3"]) if not (dataset.custom2 or "").strip(): dataset.custom2 = _dump_video_manifest(_videos_base_rel(videos_dir_rel), [], max_len=1000) dataset.save(update_fields=["custom2"]) return videos_dir_rel def _extract_sample_index(saved_filename: str) -> int | None: 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 _next_sample_index(dataset_id: str) -> int: max_saved_filename = AiDatasetSample.objects.filter(dataset_id=dataset_id).aggregate( max_saved_filename=Max("saved_filename") )["max_saved_filename"] current_index = _extract_sample_index(max_saved_filename) return (current_index or 0) + 1 def _generate_images_from_videos(dataset_id: str, frame_interval: int) -> None: close_old_connections() try: import cv2 except Exception as e: print(f"[video->images] cv2 not available: {e}") close_old_connections() return dataset = AiDataset.objects.filter(id=dataset_id).first() if not dataset: close_old_connections() return if dataset.custom2 and not str(dataset.custom2).strip().startswith(("[", "{")) and not dataset.custom3: dataset.custom3 = str(dataset.custom2).strip() dataset.custom2 = "" dataset.save(update_fields=["custom2", "custom3"]) videos_dir_rel = _get_videos_dir_rel(dataset) if not videos_dir_rel: print(f"[video->images] dataset_id={dataset_id} no videos dir") close_old_connections() return interval = int(frame_interval) if int(frame_interval) > 0 else 1 base_url = _videos_base_rel(videos_dir_rel) _, raw_items = _load_video_manifest(dataset.custom2 or "") videos: List[dict] = [] for it in raw_items: if not isinstance(it, dict): continue filename = it.get("f") or it.get("filename") if not filename and it.get("url"): filename = os.path.basename(str(it.get("url") or "")) if not filename and it.get("rel_path"): filename = os.path.basename(str(it.get("rel_path") or "")) filename = str(filename or "").strip() if not filename: continue g = it.get("g") if g is None: g = it.get("generated") videos.append( { "f": filename, "g": 1 if bool(g) else 0, "t": it.get("t") or it.get("generated_at"), "i": it.get("i") or it.get("frame_interval"), "n": it.get("n") or it.get("generated_images"), } ) pending = [v for v in videos if int(v.get("g") or 0) == 0] if not pending: print(f"[video->images] dataset_id={dataset_id} no pending videos") close_old_connections() return total_written = 0 start = time.time() now_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S") for v in pending: filename = str(v.get("f") or "").strip() if not filename: continue video_object_name = f"{str(videos_dir_rel).strip('/').replace('\\', '/')}/{filename}".strip("/") if not minio_storage.object_exists("dataset", video_object_name): print(f"[video->images] missing object: {video_object_name}") continue try: video_bytes = minio_storage.download_bytes("dataset", video_object_name) result = extract_video_frames_to_dataset_result( dataset=dataset, video_bytes=video_bytes, original_filename=filename, interval_seconds=interval, source_marker_prefix=f"video:{dataset_id}:{filename}", skip_existing=True, ) except Exception as exc: print(f"[video->images] process failed: {video_object_name}, error={exc}") continue written_for_video = len(result.get("imported") or []) total_written += written_for_video v["g"] = 1 v["t"] = now_str v["i"] = interval v["n"] = int(written_for_video) total_count, annotated_count = refresh_dataset_sample_counters(dataset_id) dataset.dataset_count = total_count dataset.annotated_count = annotated_count dataset.custom2 = _dump_video_manifest(base_url, videos, max_len=1000) dataset.save(update_fields=["dataset_count", "annotated_count", "custom2"]) print(f"[video->images] dataset_id={dataset_id} frames={total_written} elapsed={round(time.time()-start,3)}s") close_old_connections() @api_view(["POST"]) @permission_classes([IsAuthenticated]) def generate_dataset_images(request): data = request.data dataset_id = data.get("id") or data.get("dataset_id") if not dataset_id: return Response({"error": "缺少数据集 id"}, status=status.HTTP_400_BAD_REQUEST) dataset = get_object_or_404(AiDataset, id=dataset_id) if dataset.dataset_type != "03": return Response({"error": "当前数据集不是视频类型"}, status=status.HTTP_400_BAD_REQUEST) _, videos = _load_video_manifest(dataset.custom2 or "") pending = [v for v in videos if not bool(v.get("g") or v.get("generated"))] if not pending: return Response({"message": "没有需要生成图片的视频", "dataset_id": dataset_id, "pending": 0}) frame_interval = int(getattr(settings, "VIDEO_FRAME_INTERVAL", 30)) t = threading.Thread( target=_generate_images_from_videos, kwargs={"dataset_id": str(dataset_id), "frame_interval": int(frame_interval)}, daemon=True, ) t.start() return Response( { "message": "生成图片任务已后台启动", "dataset_id": dataset_id, "frame_interval": frame_interval, "pending": len(pending), "thread_name": t.name, } ) @api_view(["POST"]) @parser_classes([MultiPartParser]) @transaction.atomic def upload_dataset_file(request): data = request.data dataset = AiDataset.objects.get(id=data.get("id")) uploadfiles = request.FILES.getlist("uploadfiles[]") dataset_type = str(data.get("dataset_type") or dataset.dataset_type or "") if "uploadxml" not in data: imported_count = 0 uploadfilenames = [] interval_seconds = int(data.get("interval_seconds") or 1) max_frames = int(data.get("max_frames") or 200) for uploadfile in uploadfiles: filename = uploadfile.name uploadfilenames.append(filename) ext = os.path.splitext(filename)[1].lower() file_bytes = uploadfile.read() uploadfile.seek(0) if ext == ".zip": imported = import_zip_bytes_to_dataset(dataset, file_bytes) imported_count += len(imported) continue if dataset_type == "02" and ext in VIDEO_EXTENSIONS: imported = extract_video_frames_to_dataset( dataset, file_bytes, filename, interval_seconds=interval_seconds, max_frames=max_frames, ) imported_count += len(imported) continue if dataset_type_accepts_extension(dataset_type, ext): save_bytes_to_dataset( dataset=dataset, original_filename=filename, data=file_bytes, content_type=uploadfile.content_type, ) imported_count += 1 total_count, annotated_count = refresh_dataset_sample_counters(data.get("id")) dataset.dataset_count = total_count dataset.annotated_count = annotated_count existing_original_names = dataset.original_file_name joined_upload_names = ",".join(uploadfilenames) dataset.original_file_name = ( f"{existing_original_names},{joined_upload_names}" if existing_original_names and joined_upload_names else (existing_original_names or joined_upload_names) ) dataset.save(update_fields=["dataset_count", "annotated_count", "original_file_name"]) return Response({"message": f"文件上传成功，共导入了{imported_count}条数据！"}) uploadfiles = request.FILES.getlist("uploadfiles[]") dataset_type = str(data.get("dataset_type") or dataset.dataset_type or "") uploadfilenames = [] video_items: List[dict] = [] videos_dir_rel = "" videos_base_url = "" if dataset_type == "03": if dataset.custom2 and not str(dataset.custom2).strip().startswith(("[", "{")) and not dataset.custom3: dataset.custom3 = str(dataset.custom2).strip() dataset.custom2 = "" dataset.save(update_fields=["custom2", "custom3"]) videos_dir_rel = _ensure_video_storage(dataset) videos_base_url = _videos_base_rel(videos_dir_rel) if videos_dir_rel else "" base_existing, raw_items = _load_video_manifest(dataset.custom2 or "") if base_existing: videos_base_url = _strip_nginx_image_prefix(base_existing) for it in raw_items: if not isinstance(it, dict): continue filename = it.get("f") or it.get("filename") if not filename and it.get("url"): filename = os.path.basename(str(it.get("url") or "")) if not filename and it.get("rel_path"): filename = os.path.basename(str(it.get("rel_path") or "")) filename = str(filename or "").strip() if not filename: continue g = it.get("g") if g is None: g = it.get("generated") video_items.append({"f": filename, "g": 1 if bool(g) else 0, "t": it.get("t") or it.get("generated_at")}) def save_single_xml_file(uploadfile): file_content = uploadfile.read() uploadfile.seek(0) try: image_name, result_list = _parse_voc_xml_annotations(file_content, dataset.dataset_labels) except Exception as exc: return { "ok": False, "name": uploadfile.name, "reason": f"XML解析失败: {exc}", } dataset_sample = _find_dataset_sample_for_xml(data.get("id"), image_name, uploadfile.name) if dataset_sample is None: return { "ok": False, "name": uploadfile.name, "reason": "未找到对应样本", } annotation_json = json.dumps(result_list, ensure_ascii=False) annotation_fields = split_dataset_sample_annotation_content(annotation_json) upload_sidecar_bytes( dataset_sample.saved_path, dataset_sample.saved_filename, ".xml", file_content, "application/xml", ) AiDatasetSample.objects.filter(id=dataset_sample.id).update( rectangle=annotation_fields["rectangle"], polygon=annotation_fields["polygon"], status="02", ) return { "ok": True, "name": uploadfile.name, "sample_id": dataset_sample.id, } def process_compressed_file(sourcename, source): if not source: return source_name = os.path.basename(sourcename) source_ext = os.path.splitext(source_name)[1].lower() if dataset_type == "03" and source_ext in VIDEO_EXTENSIONS and videos_dir_rel: object_name = f"{str(videos_dir_rel).strip('/').replace('\\', '/')}/{source_name}".strip("/") content_type = "video/mp4" if source_ext == ".mp4" else None minio_storage.upload_bytes("dataset", object_name, source.read(), content_type=content_type) video_items.append({"f": source_name, "g": 0}) return save_bytes_to_dataset( dataset=dataset, original_filename=source_name, data=source.read(), ) allowed_file_types = ["jpg", "png", "bmp", "mp3", "mp4", "soud", "txt", "csv"] file_count = 0 if "uploadxml" in data: failures = [] updated_ids = [] for uploadfile in uploadfiles: resultxml = save_single_xml_file(uploadfile) if resultxml and resultxml.get("ok"): file_count += 1 updated_ids.append(resultxml.get("sample_id")) uploadfilenames.append(uploadfile.name) else: failures.append(resultxml or {"name": uploadfile.name, "reason": "处理失败"}) response_message = f"xml文件上传完成，成功更新{file_count}条样本" if failures: response_message += f"，失败{len(failures)}条" return Response( { "message": response_message, "updated_count": file_count, "updated_sample_ids": updated_ids, "failed_count": len(failures), "failures": failures, } ) for uploadfile in uploadfiles: uploadfilenames.append(uploadfile.name) file_bytes = uploadfile.read() uploadfile.seek(0) file_extension = os.path.splitext(uploadfile.name)[1].lower() if dataset_type != "01": if file_extension == ".zip": with zipfile.ZipFile(io.BytesIO(file_bytes), "r") as zip_file: for file_info in zip_file.infolist(): with zip_file.open(file_info.filename, "r", force_zip64=True) as source: decoded_filename = source.name.encode("cp437").decode("gbk") file_extension = decoded_filename.split(".")[-1].lower() if file_extension in allowed_file_types: file_count = file_count + 1 process_compressed_file(decoded_filename, source) elif file_extension == ".tar": with tarfile.open(fileobj=io.BytesIO(file_bytes), mode="r:*") as tar_file: for member in tar_file.getmembers(): extracted = tar_file.extractfile(member) if extracted is None: continue with extracted as source: file_extension = member.name.split(".")[-1].lower() if file_extension in allowed_file_types: file_count = file_count + 1 process_compressed_file(member.name, source) else: if file_extension.lstrip(".") in allowed_file_types: file_count = file_count + 1 if dataset_type == "03" and file_extension in VIDEO_EXTENSIONS and videos_dir_rel: object_name = f"{str(videos_dir_rel).strip('/').replace('\\', '/')}/{uploadfile.name}".strip("/") minio_storage.upload_bytes( "dataset", object_name, file_bytes, content_type=uploadfile.content_type or "application/octet-stream", ) video_items.append({"f": uploadfile.name, "g": 0}) else: save_bytes_to_dataset( dataset=dataset, original_filename=uploadfile.name, data=file_bytes, content_type=uploadfile.content_type, ) total_count, annotated_count = refresh_dataset_sample_counters(data.get("id")) dataset.dataset_count = total_count dataset.annotated_count = annotated_count existing_original_names = dataset.original_file_name joined_upload_names = ",".join(uploadfilenames) dataset.original_file_name = ( f"{existing_original_names},{joined_upload_names}" if existing_original_names else joined_upload_names ) if dataset_type == "03": dataset.custom2 = _dump_video_manifest(_strip_nginx_image_prefix(videos_base_url), video_items, max_len=1000) dataset.save() return Response({"message": f"文件上传成功，共上传了{file_count}条数据！"}) @api_view(["POST"]) @permission_classes([IsAuthenticated]) @transaction.atomic def create_dataset(request): data = request.data.copy() data["id"] = str(uuid.uuid4()) data["create_time"] = datetime.now() data["status"] = "01" data["creator"] = request.user.username or "creator" new_dataset_code = _generate_next_dataset_code() data["dataset_code"] = new_dataset_code dataset_paths = _build_dataset_paths(new_dataset_code) data["original_file_path"] = data.get("original_file_path") or dataset_paths["original_file_path"] data["dataset_path"] = data.get("dataset_path") or dataset_paths["dataset_path"] data["custom1"] = data.get("custom1") or dataset_paths["custom1"] if data.get("dataset_type") == "03": videos_rel = data.get("custom3") or dataset_paths["custom3"] data["custom3"] = videos_rel data["custom2"] = _dump_video_manifest(_videos_base_rel(videos_rel), [], max_len=1000) serializer = AiDatasetSerializer(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(["GET"]) @permission_classes([IsAuthenticated]) def read_dataset(request): dataset_id = request.GET.get("id", None) dataset = get_object_or_404(AiDataset, id=dataset_id) serializer = AiDatasetSerializer(dataset) STATUSES = ["00", "01", "02"] sample_stats = {} for status_code in STATUSES: sample_stats[f"count_{status_code}"] = Sum( Case(When(status=status_code, then=1), default=0, output_field=IntegerField()) ) sample_stats["total"] = Count("id") dataset_status_counts = ( AiDatasetSample.objects.values("dataset_id").annotate(**sample_stats).filter(dataset_id=dataset_id) ) serializer_data = serializer.data serializer_data["dataset_status_counts"] = dataset_status_counts return Response(serializer_data) @api_view(["PUT"]) @permission_classes([IsAuthenticated]) def update_dataset(request): data = request.data.copy() dataset = get_object_or_404(AiDataset, id=data.get("id")) data.pop("dataset_code", None) serializer = AiDatasetSerializer(instance=dataset, 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]) @transaction.atomic def delete_dataset(request): dataset_id = request.GET.get("id", None) if not dataset_id: return Response({"error": "数据集 ID 不能为空"}, status=status.HTTP_400_BAD_REQUEST) task_count = AiAnnotateTask.objects.filter(dataset_id=dataset_id).count() if task_count > 0: return Response( {"warn": "该数据集已经创建了样本标注任务，请先删除样本标注任务！"}, status=status.HTTP_400_BAD_REQUEST, ) train_count = AiAlgorithmTrainRecords.objects.filter(link_datasets__contains=dataset_id).count() if train_count > 0: return Response( {"warn": "该数据集已经创建了算法训练任务，请先删除算法训练任务！"}, status=status.HTTP_400_BAD_REQUEST, ) try: dataset = get_object_or_404(AiDataset, id=dataset_id) AiDatasetSample.objects.filter(dataset_id=dataset_id).delete() try: objects = minio_storage.list_objects("dataset", prefix=dataset.dataset_path) object_names = [obj["name"] for obj in objects if obj.get("name")] if object_names: minio_storage.delete_objects("dataset", object_names) except Exception: pass dataset.delete() return Response({"message": "数据集删除成功"}, status=status.HTTP_200_OK) except Exception as e: return Response( {"error": f"删除数据集失败: {str(e)}"}, status=status.HTTP_500_INTERNAL_SERVER_ERROR ) @api_view(["POST"]) @permission_classes([IsAuthenticated]) def delete_dataset_updatefile(request): data = request.data dataset = get_object_or_404(AiDataset, id=data.get("id")) serializer = AiDatasetSerializer(instance=dataset, 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(["GET"]) @permission_classes([IsAuthenticated]) def search_datasets(request): dataset_code = request.GET.get("dataset_code", None) dataset_name = request.GET.get("dataset_name", None) features = request.GET.get("features", None) dataset_type = request.GET.get("dataset_type", None) creator = request.GET.get("creator", None) page_size = request.GET.get("page_size", None) filter_kwargs = {} if dataset_code: filter_kwargs["dataset_code__contains"] = dataset_code if dataset_name: filter_kwargs["dataset_name__contains"] = dataset_name if dataset_type: filter_kwargs["dataset_type"] = dataset_type if features: filter_kwargs["features__contains"] = features if creator: filter_kwargs["creator"] = creator datasets = AiDataset.objects.filter(**filter_kwargs).order_by("dataset_code") paginator = PageNumberPagination() paginator.page_size = page_size or 10 result_page = paginator.paginate_queryset(datasets, request) serializer = AiDatasetSerializer(result_page, many=True) return paginator.get_paginated_response(serializer.data) @api_view(["GET"]) @permission_classes([IsAuthenticated]) def get_dataset_list(request): datasets = AiDataset.objects.all().order_by("dataset_code") serializer = AiDatasetSerializer(datasets, many=True) return Response(serializer.data) @api_view(["POST"]) @permission_classes([IsAuthenticated]) def import_dataset_from_temp_file(request): data = request.data dataset_id = data.get("dataset_id") temp_file_ids = _normalize_temp_file_ids(data) if not dataset_id or not temp_file_ids: return Response({"error": "缺少 dataset_id 或 temp_file_id(s)"}, status=status.HTTP_400_BAD_REQUEST) dataset = get_object_or_404(AiDataset, id=dataset_id) interval_seconds = int(data.get("interval_seconds") or 1) max_frames = int(data.get("max_frames") or 200) if len(temp_file_ids) == 1: temp_file = get_object_or_404(AiTempFile, id=temp_file_ids[0]) status_code, payload = _import_single_temp_file_to_dataset(dataset, temp_file, interval_seconds, max_frames) return Response(payload, status=status_code) summary_results = [] total_imported_count = 0 total_duplicate_count = 0 total_failed_count = 0 success_temp_file_count = 0 duplicate_temp_file_count = 0 failed_temp_file_count = 0 temp_files = {str(item.id): item for item in AiTempFile.objects.filter(id__in=temp_file_ids)} for temp_file_id in temp_file_ids: temp_file = temp_files.get(str(temp_file_id)) if not temp_file: payload = {"error": f"临时文件不存在: {temp_file_id}"} status_code = status.HTTP_404_NOT_FOUND else: try: status_code, payload = _import_single_temp_file_to_dataset(dataset, temp_file, interval_seconds, max_frames) except Exception as exc: logger.exception("batch temp file import failed: dataset=%s temp_file=%s", dataset.id, temp_file_id) status_code = status.HTTP_500_INTERNAL_SERVER_ERROR payload = {"error": f"导入失败：{exc}"} imported_count = int(payload.get("imported_count") or 0) duplicate_count = int(payload.get("duplicate_count") or 0) failed_count = int(payload.get("failed_count") or 0) result_status = "failed" if status_code == status.HTTP_200_OK and failed_count > 0: result_status = "partial" elif status_code == status.HTTP_200_OK: result_status = "success" elif status_code == status.HTTP_409_CONFLICT: result_status = "duplicate" if result_status in ("success", "partial"): success_temp_file_count += 1 elif result_status == "duplicate": duplicate_temp_file_count += 1 else: failed_temp_file_count += 1 total_imported_count += imported_count total_duplicate_count += duplicate_count total_failed_count += failed_count summary_results.append( { "temp_file_id": temp_file_id, "filename": getattr(temp_file, "original_name", "") if temp_file else "", "status": result_status, **payload, } ) if success_temp_file_count > 0: message = ( f"同步导入完成，成功处理{success_temp_file_count}个临时文件，" f"新增样本{total_imported_count}条，跳过重复{total_duplicate_count}条，失败{total_failed_count}条" ) response_status = status.HTTP_200_OK elif duplicate_temp_file_count > 0 and failed_temp_file_count == 0: message = "所选临时文件均已导入当前数据集，本次未新增样本" response_status = status.HTTP_200_OK else: message = "同步导入失败，请检查临时文件类型与数据集类型是否匹配" response_status = status.HTTP_400_BAD_REQUEST return Response( { "message": message, "dataset_id": dataset.id, "temp_file_count": len(temp_file_ids), "success_temp_file_count": success_temp_file_count, "duplicate_temp_file_count": duplicate_temp_file_count, "failed_temp_file_count": failed_temp_file_count, "imported_count": total_imported_count, "duplicate_count": total_duplicate_count, "failed_count": total_failed_count, "results": summary_results, }, status=response_status, )