1102 lines
43 KiB
Python
1102 lines
43 KiB
Python
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,
|
||
)
|