503 lines
17 KiB
Python
503 lines
17 KiB
Python
import io
|
|
import mimetypes
|
|
import os
|
|
import uuid
|
|
import zipfile
|
|
from datetime import datetime
|
|
from tempfile import NamedTemporaryFile
|
|
|
|
from django.conf import settings
|
|
from django.db.models import Max
|
|
from PIL import Image, UnidentifiedImageError
|
|
|
|
from apps.common.minio_client import minio_storage
|
|
from apps.common.utils.annotation_store import split_dataset_sample_annotation_content
|
|
from apps.core.models import AiDatasetSample
|
|
|
|
try:
|
|
import cv2
|
|
except Exception:
|
|
cv2 = None
|
|
|
|
|
|
DATASET_BUCKET = "dataset"
|
|
TEMPFILE_BUCKET = "tempfile"
|
|
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".gif", ".webp", ".tif", ".tiff"}
|
|
VIDEO_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".wmv", ".flv", ".mpeg", ".mpg", ".webm"}
|
|
AUDIO_EXTENSIONS = {".mp3", ".wav", ".aac", ".m4a", ".flac", ".ogg"}
|
|
TEXT_EXTENSIONS = {".txt", ".csv", ".json", ".xml"}
|
|
|
|
|
|
def _local_sample_path(train_path: str,saved_path: str, saved_filename: str) -> str:
|
|
return os.path.join(train_path, saved_path, saved_filename).replace("\\", "/")
|
|
|
|
|
|
def _sample_object_name(saved_path: str, saved_filename: str) -> str:
|
|
return f"{str(saved_path or '').strip('/').replace('\\', '/')}/{saved_filename}".strip("/")
|
|
|
|
|
|
def _ensure_local_dir(saved_path: str) -> str:
|
|
local_dir = os.path.join(settings.FILESPACE_ROOT_PATH, saved_path).replace("\\", "/")
|
|
os.makedirs(local_dir, exist_ok=True)
|
|
return local_dir
|
|
|
|
|
|
def _write_local_file(saved_path: str, saved_filename: str, data: bytes) -> str:
|
|
local_dir = _ensure_local_dir(saved_path)
|
|
full_path = os.path.join(local_dir, saved_filename).replace("\\", "/")
|
|
with open(full_path, "wb") as f:
|
|
f.write(data)
|
|
return full_path
|
|
|
|
|
|
def _sidecar_filename(saved_filename: str, suffix: str) -> str:
|
|
stem = os.path.splitext(saved_filename)[0]
|
|
normalized_suffix = suffix if str(suffix).startswith(".") else f".{suffix}"
|
|
return f"{stem}{normalized_suffix}"
|
|
|
|
|
|
def sample_sidecar_names(saved_filename: str):
|
|
return [
|
|
_sidecar_filename(saved_filename, ".xml"),
|
|
_sidecar_filename(saved_filename, ".json"),
|
|
]
|
|
|
|
|
|
def _normalize_filesize(size_bytes: int) -> str:
|
|
try:
|
|
return str(min(max(int(size_bytes or 0), 0), 9999999999))
|
|
except Exception:
|
|
return "0"
|
|
|
|
|
|
def _image_wh_from_bytes(data: bytes):
|
|
try:
|
|
with Image.open(io.BytesIO(data)) as img:
|
|
width, height = img.size
|
|
return f"{width}x{height}"
|
|
except UnidentifiedImageError:
|
|
return None
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
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 _build_saved_filename(dataset, original_filename: str) -> str:
|
|
ext = os.path.splitext(original_filename or "")[1].lower()
|
|
index = _next_sample_index(dataset.id)
|
|
return f"{dataset.dataset_code}-{index:08}{ext}"
|
|
|
|
|
|
def detect_original_type(filename: str) -> str:
|
|
ext = os.path.splitext(filename or "")[1].lower()
|
|
if ext in IMAGE_EXTENSIONS:
|
|
return "01"
|
|
if ext in VIDEO_EXTENSIONS:
|
|
return "02"
|
|
if ext in AUDIO_EXTENSIONS:
|
|
return "03"
|
|
if ext == ".zip":
|
|
return "04"
|
|
if ext in TEXT_EXTENSIONS:
|
|
return "05"
|
|
return "99"
|
|
|
|
|
|
def dataset_type_accepts_extension(dataset_type: str, extension: str) -> bool:
|
|
ext = str(extension or "").lower()
|
|
if dataset_type == "02":
|
|
return ext in IMAGE_EXTENSIONS
|
|
if dataset_type == "03":
|
|
return ext in VIDEO_EXTENSIONS
|
|
if dataset_type == "04":
|
|
return ext in AUDIO_EXTENSIONS
|
|
if dataset_type == "05":
|
|
return ext in TEXT_EXTENSIONS
|
|
return ext in IMAGE_EXTENSIONS | VIDEO_EXTENSIONS | AUDIO_EXTENSIONS | TEXT_EXTENSIONS
|
|
|
|
|
|
def create_dataset_sample_record(
|
|
dataset,
|
|
original_filename: str,
|
|
saved_filename: str,
|
|
size_bytes: int,
|
|
annotation_content=None,
|
|
custom1=None,
|
|
custom2=None,
|
|
custom3=None,
|
|
status: str = "00",
|
|
):
|
|
annotation_fields = split_dataset_sample_annotation_content(annotation_content)
|
|
return AiDatasetSample.objects.create(
|
|
id=str(uuid.uuid4()),
|
|
dataset_id=dataset.id,
|
|
saved_path=dataset.dataset_path,
|
|
original_filename=original_filename,
|
|
saved_filename=saved_filename,
|
|
saved_filesize=_normalize_filesize(size_bytes),
|
|
rectangle=annotation_fields["rectangle"],
|
|
polygon=annotation_fields["polygon"],
|
|
status=status,
|
|
create_time=datetime.now(),
|
|
custom1=custom1,
|
|
custom2=custom2,
|
|
custom3=custom3,
|
|
)
|
|
|
|
|
|
def save_bytes_to_dataset(
|
|
dataset,
|
|
original_filename: str,
|
|
data: bytes,
|
|
content_type: str | None = None,
|
|
annotation_content=None,
|
|
custom1=None,
|
|
custom2=None,
|
|
custom3=None,
|
|
):
|
|
saved_filename = _build_saved_filename(dataset, original_filename)
|
|
object_name = _sample_object_name(dataset.dataset_path, saved_filename)
|
|
actual_content_type = content_type or mimetypes.guess_type(original_filename)[0] or "application/octet-stream"
|
|
minio_storage.upload_bytes(DATASET_BUCKET, object_name, data, content_type=actual_content_type)
|
|
if custom1 is None and os.path.splitext(saved_filename)[1].lower() in IMAGE_EXTENSIONS:
|
|
custom1 = _image_wh_from_bytes(data)
|
|
sample = create_dataset_sample_record(
|
|
dataset=dataset,
|
|
original_filename=os.path.basename(original_filename),
|
|
saved_filename=saved_filename,
|
|
size_bytes=len(data),
|
|
annotation_content=annotation_content,
|
|
custom1=custom1,
|
|
custom2=custom2,
|
|
custom3=custom3,
|
|
)
|
|
return sample
|
|
|
|
|
|
def save_uploaded_file_to_dataset(dataset, uploaded_file):
|
|
file_bytes = uploaded_file.read()
|
|
uploaded_file.seek(0)
|
|
return save_bytes_to_dataset(
|
|
dataset=dataset,
|
|
original_filename=uploaded_file.name,
|
|
data=file_bytes,
|
|
content_type=uploaded_file.content_type,
|
|
)
|
|
|
|
|
|
def copy_minio_object_to_dataset(
|
|
dataset,
|
|
source_bucket: str,
|
|
source_object: str,
|
|
original_filename: str,
|
|
size_bytes: int | None = None,
|
|
custom2=None,
|
|
custom3=None,
|
|
):
|
|
saved_filename = _build_saved_filename(dataset, original_filename)
|
|
target_object = _sample_object_name(dataset.dataset_path, saved_filename)
|
|
minio_storage.copy_object(source_bucket, source_object, DATASET_BUCKET, target_object)
|
|
actual_size = size_bytes
|
|
if actual_size is None:
|
|
try:
|
|
actual_size = int(minio_storage.get_object_info(DATASET_BUCKET, target_object).get("size") or 0)
|
|
except Exception:
|
|
actual_size = 0
|
|
custom1 = None
|
|
if os.path.splitext(saved_filename)[1].lower() in IMAGE_EXTENSIONS:
|
|
try:
|
|
custom1 = _image_wh_from_bytes(minio_storage.download_bytes(DATASET_BUCKET, target_object))
|
|
except Exception:
|
|
custom1 = None
|
|
return create_dataset_sample_record(
|
|
dataset=dataset,
|
|
original_filename=os.path.basename(original_filename),
|
|
saved_filename=saved_filename,
|
|
size_bytes=actual_size or 0,
|
|
custom1=custom1,
|
|
custom2=custom2,
|
|
custom3=custom3,
|
|
)
|
|
|
|
|
|
def dataset_sample_source_exists(dataset_id: str, source_marker: str | None) -> bool:
|
|
if not source_marker:
|
|
return False
|
|
return AiDatasetSample.objects.filter(dataset_id=dataset_id, custom2=source_marker).exists()
|
|
|
|
|
|
def copy_minio_object_to_dataset_result(
|
|
dataset,
|
|
source_bucket: str,
|
|
source_object: str,
|
|
original_filename: str,
|
|
size_bytes: int | None = None,
|
|
source_marker: str | None = None,
|
|
skip_existing: bool = False,
|
|
):
|
|
if skip_existing and dataset_sample_source_exists(dataset.id, source_marker):
|
|
return {
|
|
"imported": [],
|
|
"skipped": [{"name": os.path.basename(original_filename), "reason": "duplicate_source"}],
|
|
"failed": [],
|
|
}
|
|
sample = copy_minio_object_to_dataset(
|
|
dataset=dataset,
|
|
source_bucket=source_bucket,
|
|
source_object=source_object,
|
|
original_filename=original_filename,
|
|
size_bytes=size_bytes,
|
|
custom2=source_marker,
|
|
)
|
|
return {"imported": [sample], "skipped": [], "failed": []}
|
|
|
|
|
|
def import_zip_bytes_to_dataset_result(
|
|
dataset,
|
|
zip_bytes: bytes,
|
|
source_marker_prefix: str | None = None,
|
|
skip_existing: bool = False,
|
|
):
|
|
imported = []
|
|
skipped = []
|
|
failed = []
|
|
with zipfile.ZipFile(io.BytesIO(zip_bytes), "r") as zip_file:
|
|
for member in zip_file.infolist():
|
|
if member.is_dir():
|
|
continue
|
|
member_name = member.filename.replace("\\", "/")
|
|
filename = os.path.basename(member_name)
|
|
if not filename or filename.startswith("."):
|
|
continue
|
|
ext = os.path.splitext(filename)[1].lower()
|
|
if not dataset_type_accepts_extension(dataset.dataset_type, ext):
|
|
failed.append({"name": filename, "reason": "type_mismatch"})
|
|
continue
|
|
source_marker = f"{source_marker_prefix}:zip:{member_name}" if source_marker_prefix else None
|
|
if skip_existing and dataset_sample_source_exists(dataset.id, source_marker):
|
|
skipped.append({"name": filename, "reason": "duplicate_source"})
|
|
continue
|
|
try:
|
|
with zip_file.open(member, "r") as fp:
|
|
data = fp.read()
|
|
imported.append(save_bytes_to_dataset(dataset, filename, data, custom2=source_marker))
|
|
except Exception as exc:
|
|
failed.append({"name": filename, "reason": str(exc)})
|
|
return {"imported": imported, "skipped": skipped, "failed": failed}
|
|
|
|
|
|
def import_zip_bytes_to_dataset(dataset, zip_bytes: bytes, source_marker_prefix: str | None = None, skip_existing: bool = False):
|
|
result = import_zip_bytes_to_dataset_result(
|
|
dataset=dataset,
|
|
zip_bytes=zip_bytes,
|
|
source_marker_prefix=source_marker_prefix,
|
|
skip_existing=skip_existing,
|
|
)
|
|
return result["imported"]
|
|
|
|
|
|
def extract_video_frames_to_dataset_result(
|
|
dataset,
|
|
video_bytes: bytes,
|
|
original_filename: str,
|
|
interval_seconds: int = 1,
|
|
max_frames: int = 200,
|
|
image_format: str = "jpg",
|
|
source_marker_prefix: str | None = None,
|
|
skip_existing: bool = False,
|
|
):
|
|
if cv2 is None:
|
|
raise RuntimeError("opencv-python 未安装，无法处理视频")
|
|
|
|
video_ext = os.path.splitext(original_filename or "")[1].lower() or ".mp4"
|
|
temp_path = None
|
|
cap = None
|
|
imported = []
|
|
skipped = []
|
|
failed = []
|
|
try:
|
|
with NamedTemporaryFile(delete=False, suffix=video_ext) as temp_file:
|
|
temp_file.write(video_bytes)
|
|
temp_path = temp_file.name
|
|
|
|
cap = cv2.VideoCapture(temp_path)
|
|
if not cap.isOpened():
|
|
raise RuntimeError("视频文件无法打开")
|
|
|
|
fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
|
|
step_frames = max(int(round(interval_seconds * fps)), 1)
|
|
current_frame = 0
|
|
video_stem = os.path.splitext(os.path.basename(original_filename))[0]
|
|
suffix = "jpg" if image_format == "jpeg" else image_format
|
|
encode_ext = ".jpg" if image_format in {"jpg", "jpeg"} else ".png"
|
|
|
|
while len(imported) < max_frames:
|
|
ok, frame = cap.read()
|
|
if not ok:
|
|
break
|
|
if current_frame % step_frames == 0:
|
|
source_marker = (
|
|
f"{source_marker_prefix}:frame:{current_frame:08}"
|
|
if source_marker_prefix
|
|
else None
|
|
)
|
|
if skip_existing and dataset_sample_source_exists(dataset.id, source_marker):
|
|
skipped.append(
|
|
{
|
|
"name": f"{video_stem}_{current_frame:08}.{suffix}",
|
|
"reason": "duplicate_source",
|
|
}
|
|
)
|
|
current_frame += 1
|
|
continue
|
|
ok_encode, encoded = cv2.imencode(encode_ext, frame)
|
|
if ok_encode:
|
|
frame_name = f"{video_stem}_{current_frame:08}.{suffix}"
|
|
imported.append(
|
|
save_bytes_to_dataset(
|
|
dataset,
|
|
frame_name,
|
|
encoded.tobytes(),
|
|
custom2=source_marker,
|
|
)
|
|
)
|
|
else:
|
|
failed.append(
|
|
{
|
|
"name": f"{video_stem}_{current_frame:08}.{suffix}",
|
|
"reason": "frame_encode_failed",
|
|
}
|
|
)
|
|
current_frame += 1
|
|
finally:
|
|
if cap is not None:
|
|
cap.release()
|
|
if temp_path and os.path.exists(temp_path):
|
|
os.unlink(temp_path)
|
|
return {"imported": imported, "skipped": skipped, "failed": failed}
|
|
|
|
|
|
def extract_video_frames_to_dataset(
|
|
dataset,
|
|
video_bytes: bytes,
|
|
original_filename: str,
|
|
interval_seconds: int = 1,
|
|
max_frames: int = 200,
|
|
image_format: str = "jpg",
|
|
source_marker_prefix: str | None = None,
|
|
skip_existing: bool = False,
|
|
):
|
|
result = extract_video_frames_to_dataset_result(
|
|
dataset=dataset,
|
|
video_bytes=video_bytes,
|
|
original_filename=original_filename,
|
|
interval_seconds=interval_seconds,
|
|
max_frames=max_frames,
|
|
image_format=image_format,
|
|
source_marker_prefix=source_marker_prefix,
|
|
skip_existing=skip_existing,
|
|
)
|
|
return result["imported"]
|
|
|
|
|
|
def ensure_sample_local_copy(train_path: str, saved_path: str, saved_filename: str):
|
|
local_path = _local_sample_path(train_path, saved_path, saved_filename)
|
|
if os.path.exists(local_path):
|
|
return local_path
|
|
object_name = _sample_object_name(saved_path, saved_filename)
|
|
_ensure_local_dir(saved_path)
|
|
minio_storage.download_file(DATASET_BUCKET, object_name, local_path)
|
|
return local_path
|
|
|
|
|
|
def ensure_sidecar_local_copy(saved_path: str, saved_filename: str, suffix: str):
|
|
sidecar_name = _sidecar_filename(saved_filename, suffix)
|
|
local_path = _local_sample_path(saved_path, sidecar_name)
|
|
if os.path.exists(local_path):
|
|
return local_path
|
|
object_name = _sample_object_name(saved_path, sidecar_name)
|
|
if not minio_storage.object_exists(DATASET_BUCKET, object_name):
|
|
return local_path
|
|
_ensure_local_dir(saved_path)
|
|
minio_storage.download_file(DATASET_BUCKET, object_name, local_path)
|
|
return local_path
|
|
|
|
|
|
def download_sidecar_bytes(saved_path: str, saved_filename: str, suffix: str):
|
|
sidecar_name = _sidecar_filename(saved_filename, suffix)
|
|
object_name = _sample_object_name(saved_path, sidecar_name)
|
|
try:
|
|
if minio_storage.object_exists(DATASET_BUCKET, object_name):
|
|
return minio_storage.download_bytes(DATASET_BUCKET, object_name)
|
|
except Exception:
|
|
pass
|
|
|
|
local_path = _local_sample_path(saved_path, sidecar_name)
|
|
if os.path.exists(local_path) and os.path.isfile(local_path):
|
|
with open(local_path, "rb") as f:
|
|
return f.read()
|
|
return None
|
|
|
|
|
|
def download_sidecar_text(saved_path: str, saved_filename: str, suffix: str, encoding: str = "utf-8"):
|
|
data = download_sidecar_bytes(saved_path, saved_filename, suffix)
|
|
if data is None:
|
|
return None
|
|
return data.decode(encoding, errors="replace")
|
|
|
|
|
|
def upload_sidecar_bytes(saved_path: str, saved_filename: str, suffix: str, data: bytes, content_type: str | None = None):
|
|
sidecar_name = _sidecar_filename(saved_filename, suffix)
|
|
object_name = _sample_object_name(saved_path, sidecar_name)
|
|
actual_content_type = content_type or mimetypes.guess_type(sidecar_name)[0] or "application/octet-stream"
|
|
minio_storage.upload_bytes(DATASET_BUCKET, object_name, data, content_type=actual_content_type)
|
|
return sidecar_name
|
|
|
|
|
|
def upload_sidecar_text(saved_path: str, saved_filename: str, suffix: str, text: str, content_type: str | None = None):
|
|
actual_content_type = content_type
|
|
if actual_content_type is None and str(suffix).lower().endswith("json"):
|
|
actual_content_type = "application/json"
|
|
return upload_sidecar_bytes(
|
|
saved_path=saved_path,
|
|
saved_filename=saved_filename,
|
|
suffix=suffix,
|
|
data=(text or "").encode("utf-8"),
|
|
content_type=actual_content_type,
|
|
)
|
|
|
|
|
|
def delete_sample_storage(saved_path: str, saved_filename: str):
|
|
candidates = [saved_filename] + sample_sidecar_names(saved_filename)
|
|
failed_objects = []
|
|
for filename in candidates:
|
|
object_name = _sample_object_name(saved_path, filename)
|
|
try:
|
|
if minio_storage.object_exists(DATASET_BUCKET, object_name):
|
|
minio_storage.delete_object(DATASET_BUCKET, object_name)
|
|
except Exception:
|
|
failed_objects.append(object_name)
|
|
|
|
local_path = os.path.join(settings.FILESPACE_ROOT_PATH, saved_path, filename).replace("\\", "/")
|
|
if os.path.exists(local_path) and os.path.isfile(local_path):
|
|
try:
|
|
os.remove(local_path)
|
|
except Exception:
|
|
pass
|
|
return failed_objects
|