Refine dataset mask upload workflow

This commit is contained in:
2026-07-01 11:07:32 +08:00
parent 1e8339e237
commit afd6bd39ec
4 changed files with 206 additions and 102 deletions

View File

@@ -3,6 +3,8 @@ from __future__ import annotations
import json
import re
import shutil
import tarfile
import zipfile
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterable
@@ -14,6 +16,12 @@ from ...config import settings
DATASET_KINDS = ("images", "labels", "masks")
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}
LABEL_EXTS = {".txt", ".json", ".yaml", ".yml"}
ARCHIVE_SUFFIXES = (".zip", ".tar", ".tar.gz", ".tgz")
ARCHIVE_ALIASES = {
"images": {"image", "images", "img", "imgs", "ori", "original", "originals"},
"masks": {"mask", "masks", "label", "labels", "gt", "annotation", "annotations"},
"labels": {"label", "labels", "txt", "annotation", "annotations"},
}
def uploads_root() -> Path:
@@ -36,6 +44,68 @@ def safe_filename(value: str | None) -> str:
return stem
def is_archive_name(value: str | None) -> bool:
name = (value or "").lower()
return any(name.endswith(suffix) for suffix in ARCHIVE_SUFFIXES)
def unique_path(path: Path) -> Path:
if not path.exists():
return path
stem = path.stem
suffix = path.suffix
counter = 1
while True:
candidate = path.with_name(f"{stem}_{counter}{suffix}")
if not candidate.exists():
return candidate
counter += 1
def safe_archive_member(raw_name: str, kind: str) -> Path | None:
normalized = raw_name.replace("\\", "/")
if not normalized or normalized.endswith("/"):
return None
raw_parts = [part for part in normalized.split("/") if part not in {"", "."}]
if not raw_parts or any(part == ".." for part in raw_parts):
raise ValueError(f"unsafe archive member path: {raw_name}")
if Path(normalized).is_absolute():
raise ValueError(f"absolute archive member path is not allowed: {raw_name}")
lower_parts = [part.lower() for part in raw_parts]
if lower_parts[0] in {"__macosx", ".ds_store"}:
return None
target_aliases = ARCHIVE_ALIASES.get(kind, {kind})
other_aliases = set().union(*(aliases for item, aliases in ARCHIVE_ALIASES.items() if item != kind))
target_index = next((index for index, part in enumerate(lower_parts[:-1]) if part in target_aliases), None)
other_index = next((index for index, part in enumerate(lower_parts[:-1]) if part in other_aliases), None)
if target_index is not None:
raw_parts = raw_parts[target_index + 1 :]
elif other_index is not None:
return None
if not raw_parts:
return None
safe_parts = [slugify(part) for part in raw_parts[:-1]]
filename = safe_filename(raw_parts[-1])
return Path(*safe_parts, filename) if safe_parts else Path(filename)
def save_archive_member(target: Path, member_name: str, kind: str, source) -> dict | None:
relative = safe_archive_member(member_name, kind)
if relative is None:
return None
dst = unique_path(target / relative)
resolved_target = target.resolve()
resolved_dst = dst.resolve()
if resolved_target not in resolved_dst.parents and resolved_dst != resolved_target:
raise ValueError(f"archive member escapes target directory: {member_name}")
dst.parent.mkdir(parents=True, exist_ok=True)
with dst.open("wb") as handle:
shutil.copyfileobj(source, handle)
return {"name": dst.name, "relative_path": str(dst.relative_to(settings.project_root)), "size": dst.stat().st_size, "from_archive": True}
def dataset_dir(name: str) -> Path:
return uploads_root() / slugify(name)
@@ -283,15 +353,35 @@ async def save_upload(dataset: str, kind: str, files: list[UploadFile]) -> dict:
saved = []
for upload in files:
filename = safe_filename(upload.filename)
dst = target / filename
if dst.exists():
stem = dst.stem
suffix = dst.suffix
counter = 1
while dst.exists():
dst = target / f"{stem}_{counter}{suffix}"
counter += 1
upload.file.seek(0)
if is_archive_name(upload.filename):
archive_saved = []
if filename.lower().endswith(".zip"):
with zipfile.ZipFile(upload.file) as archive:
for info in archive.infolist():
if info.is_dir():
continue
with archive.open(info) as source:
item = save_archive_member(target, info.filename, kind, source)
if item:
archive_saved.append(item)
else:
with tarfile.open(fileobj=upload.file, mode="r:*") as archive:
for member in archive:
if not member.isfile():
continue
source = archive.extractfile(member)
if source is None:
continue
with source:
item = save_archive_member(target, member.name, kind, source)
if item:
archive_saved.append(item)
saved.extend(archive_saved)
continue
dst = unique_path(target / filename)
with dst.open("wb") as handle:
shutil.copyfileobj(upload.file, handle)
saved.append({"name": dst.name, "relative_path": str(dst.relative_to(settings.project_root)), "size": dst.stat().st_size})
saved.append({"name": dst.name, "relative_path": str(dst.relative_to(settings.project_root)), "size": dst.stat().st_size, "from_archive": False})
return {"dataset": describe_dataset(safe_name), "saved": saved}

View File

@@ -1,7 +1,12 @@
import asyncio
import io
import zipfile
import cv2
import numpy as np
from fastapi import UploadFile
from app.modules.dataset.service import create_dataset, describe_dataset, generate_yolo_dataset_yaml, validate_dataset
from app.modules.dataset.service import create_dataset, describe_dataset, generate_yolo_dataset_yaml, save_upload, validate_dataset
def test_create_dataset_layout(tmp_path, monkeypatch):
@@ -40,3 +45,33 @@ def test_validate_dataset_and_generate_yolo_yaml(tmp_path, monkeypatch):
assert generated["relative_path"] == "var/uploads/datasets/case_yolo/dataset.yaml"
assert "nc: 1" in generated["content"]
assert "0: tool" in generated["content"]
def test_upload_zip_extracts_matching_dataset_kind(tmp_path, monkeypatch):
from types import SimpleNamespace
from app.modules.dataset import service
monkeypatch.setattr(service, "settings", SimpleNamespace(project_root=tmp_path))
create_dataset("case zip", "demo")
image_zip = io.BytesIO()
with zipfile.ZipFile(image_zip, "w") as archive:
archive.writestr("bundle/images/sample.png", b"image")
archive.writestr("bundle/masks/sample.png", b"mask")
image_zip.seek(0)
image_upload = UploadFile(filename="images.zip", file=image_zip)
image_result = asyncio.run(save_upload("case_zip", "images", [image_upload]))
mask_zip = io.BytesIO()
with zipfile.ZipFile(mask_zip, "w") as archive:
archive.writestr("bundle/images/sample.png", b"image")
archive.writestr("bundle/masks/sample.png", b"mask")
mask_zip.seek(0)
mask_upload = UploadFile(filename="masks.zip", file=mask_zip)
mask_result = asyncio.run(save_upload("case_zip", "masks", [mask_upload]))
assert [item["relative_path"] for item in image_result["saved"]] == ["var/uploads/datasets/case_zip/images/sample.png"]
assert [item["relative_path"] for item in mask_result["saved"]] == ["var/uploads/datasets/case_zip/masks/sample.png"]
described = describe_dataset("case_zip")
assert described["counts"]["images"] == 1
assert described["counts"]["masks"] == 1

View File

@@ -19,12 +19,15 @@ import {
Square,
Terminal,
UploadCloud,
PackageOpen,
Wand2,
Zap
} from "lucide-react";
import "./styles.css";
const API_BASE = import.meta.env.VITE_API_BASE ?? "http://localhost:8010";
const DEFAULT_API_BASE =
typeof window !== "undefined" ? `${window.location.protocol}//${window.location.hostname}:8010` : "http://127.0.0.1:8010";
const API_BASE = (import.meta.env.VITE_API_BASE ?? DEFAULT_API_BASE).replace(/\/$/, "");
type JobProgress = {
percent: number | null;
@@ -344,7 +347,7 @@ type ModelWeightOption = {
const pages: Array<{ id: PageId; label: string; eyebrow: string; title: string; description: string; icon: React.ReactNode }> = [
{ id: "overview", label: "总览", eyebrow: "Operations Map", title: "分割平台运行总览", description: "能力矩阵、关键资产和最近产物集中看板。", icon: <Gauge size={18} /> },
{ id: "datasets", label: "数据集", eyebrow: "Dataset Bench", title: "数据集、Label、Mask 工作台", description: "按数据集管理上传、配对校验、训练数据生成和样例预览。", icon: <Boxes size={18} /> },
{ id: "datasets", label: "数据集", eyebrow: "Dataset Bench", title: "数据集Mask 工作台", description: "按数据集管理图片、masks 上传、配对校验和样例预览。", icon: <Boxes size={18} /> },
{ id: "training", label: "训练", eyebrow: "Training Queue", title: "训练任务与实时日志", description: "选择 SegModel、YOLO、MMSeg 或数据处理任务并跟踪进度。", icon: <Terminal size={18} /> },
{ id: "inference", label: "推理", eyebrow: "Model Inference", title: "选择训练模型进行图片推理", description: "从已有 best.pt/last.pt/manifest 权重中选择模型,对上传图片或数据集运行预测与热度图。", icon: <FileImage size={18} /> },
{ id: "results", label: "结果", eyebrow: "Result Studio", title: "分割结果、热度图与 Loss 曲线", description: "按模型族和产物类型浏览训练结果、推理图、热度图和 CSV 曲线。", icon: <BarChart3 size={18} /> },
@@ -358,14 +361,23 @@ function pageFromHash(): PageId {
}
async function api<T>(path: string, init?: RequestInit): Promise<T> {
const res = await fetch(`${API_BASE}${path}`, {
let res: Response;
try {
res = await fetch(`${API_BASE}${path}`, {
headers: { "Content-Type": "application/json" },
...init
});
} catch (err) {
throw new Error(`无法连接后端 ${API_BASE}${err instanceof Error ? err.message : String(err)}`);
}
if (!res.ok) throw new Error(await res.text());
return res.json();
}
function artifactUrl(path: string) {
return `${API_BASE}/api/artifacts/${path.split("/").map(encodeURIComponent).join("/")}`;
}
const defaultParams: Record<string, Record<string, unknown>> = {
"mock.echo": { message: "hello from Seg Data Server" },
"dataset.rename": { image_dir: "../DataSet_Own/ORI", label_dir: "../DataSet_Own/Label" },
@@ -395,14 +407,14 @@ const taskLabels: Record<string, string> = {
"dataset.rename": "重命名",
"dataset.to_png": "转 PNG",
"dataset.resize": "Resize",
"dataset.pair": "图片/Label 配对",
"dataset.rebuild_labels": "重建 Label",
"dataset.pair": "图片/Mask 配对",
"dataset.rebuild_labels": "重建 Mask",
"dataset.stack": "透明叠加",
"dataset.stitch": "拼接检查",
"dataset.video_frames": "视频抽帧",
"dataset.yolo_check_pairs": "YOLO 配对",
"dataset.yolo_stack": "YOLO 叠加",
"dataset.yolo_rebuild_labels": "YOLO Label",
"dataset.yolo_rebuild_labels": "YOLO Mask",
"dataset.yolo_txt_sort": "生成 TXT",
"dataset.yolo_convert_png": "批量 PNG",
"dataset.yolo_resize": "批量缩放"
@@ -547,7 +559,7 @@ function App() {
const [resultQuery, setResultQuery] = useState("");
const [selectedGpuIds, setSelectedGpuIds] = useState<number[]>([]);
const [selectedCondaEnv, setSelectedCondaEnv] = useState("auto");
const [uploadKind, setUploadKind] = useState<"images" | "labels" | "masks">("images");
const [uploadKind, setUploadKind] = useState<"images" | "masks">("images");
const [uploadFiles, setUploadFiles] = useState<FileList | null>(null);
const [agentValidation, setAgentValidation] = useState<ValidationAgentPayload | null>(null);
const [weightVerification, setWeightVerification] = useState<WeightVerifyPayload | null>(null);
@@ -665,6 +677,9 @@ function App() {
}, [results, weightManifest]);
const activePage = pages.find((item) => item.id === page) ?? pages[0];
const inferenceSource = inferenceSourcePath.trim() || selectedDataset?.absolute_layout?.images || "";
const uploadAccept = uploadKind === "images"
? "image/*,.zip,.tar,.tar.gz,.tgz"
: "image/*,.txt,.json,.yaml,.yml,.zip,.tar,.tar.gz,.tgz";
const inferenceOutputs = useMemo(() => {
const predictions = results.filter((item) => item.role === "segmentation" && item.previewable).slice(0, 8);
const heatmaps = results.filter((item) => item.role === "heatmap" && item.previewable).slice(0, 8);
@@ -730,12 +745,13 @@ function App() {
const base = { ...(catalog?.task_defaults?.[next] ?? defaultParams[next] ?? {}) };
const layout = selectedDataset?.absolute_layout;
if (!layout) return base;
const maskDir = layout.masks || layout.labels;
const resultDir = `${layout.images.replace(/\/images$/, "")}/results/${next.replace(".", "_")}`;
if (["dataset.pair", "dataset.resize", "dataset.stack", "dataset.stitch", "dataset.yolo_check_pairs", "dataset.yolo_stack"].includes(next)) {
return { ...base, image_dir: layout.images, label_dir: layout.labels, result_dir: resultDir };
return { ...base, image_dir: layout.images, label_dir: maskDir, result_dir: resultDir };
}
if (["dataset.rebuild_labels", "dataset.yolo_rebuild_labels", "dataset.yolo_txt_sort"].includes(next)) {
return { ...base, label_dir: layout.labels, folder: layout.labels };
return { ...base, label_dir: maskDir, folder: maskDir };
}
if (["dataset.to_png", "dataset.yolo_convert_png", "dataset.yolo_resize"].includes(next)) {
return { ...base, input_dir: layout.images, output_dir: resultDir, folder: layout.images };
@@ -861,10 +877,15 @@ function App() {
try {
const body = new FormData();
Array.from(uploadFiles).forEach((file) => body.append("files", file));
const res = await fetch(`${API_BASE}/api/datasets/${encodeURIComponent(datasetName)}/upload/${uploadKind}`, {
let res: Response;
try {
res = await fetch(`${API_BASE}/api/datasets/${encodeURIComponent(datasetName)}/upload/${uploadKind}`, {
method: "POST",
body
});
} catch (err) {
throw new Error(`上传失败,无法连接后端 ${API_BASE}${err instanceof Error ? err.message : String(err)}`);
}
if (!res.ok) throw new Error(await res.text());
await refresh();
} finally {
@@ -1210,7 +1231,7 @@ function App() {
<div className="panelHead">
<div>
<p className="eyebrow">Dataset Bench</p>
<h2>LabelMask </h2>
<h2> Masks </h2>
</div>
<Database size={22} />
</div>
@@ -1224,16 +1245,17 @@ function App() {
<input value={datasetDescription} onChange={(event) => setDatasetDescription(event.target.value)} />
</label>
<div className="segmented">
{(["images", "labels", "masks"] as const).map((kind) => (
{(["images", "masks"] as const).map((kind) => (
<button key={kind} className={uploadKind === kind ? "active" : ""} onClick={() => setUploadKind(kind)}>
{kind}
</button>
))}
</div>
<label className="drop">
<UploadCloud size={24} />
<span>{uploadFiles?.length ? `${uploadFiles.length} files selected` : "选择图片、label 或 mask 文件"}</span>
<input multiple type="file" accept="image/*,.txt,.json,.yaml,.yml" onChange={(event) => setUploadFiles(event.target.files)} />
{uploadFiles?.length ? <PackageOpen size={24} /> : <UploadCloud size={24} />}
<span>{uploadFiles?.length ? `${uploadFiles.length} 个文件待上传` : `选择 ${uploadKind} 图片或压缩包`}</span>
<small>.png/.jpg/.tif/.zip/.tar.gz </small>
<input multiple type="file" accept={uploadAccept} onChange={(event) => setUploadFiles(event.target.files)} />
</label>
<div className="buttonRow">
<button className="primary" disabled={busy} onClick={createDataset}><Boxes size={17} /></button>
@@ -1292,16 +1314,17 @@ function App() {
>
<div className="datasetCardHead">
<strong>{dataset.name}</strong>
<span>{dataset.counts.images} image · {dataset.counts.labels} label · {dataset.counts.masks} mask</span>
<span>{dataset.counts.images} images · {dataset.counts.masks} masks</span>
</div>
<div className="readinessLine">
<StatusPill status={datasetValidations[dataset.name]?.ready.yolo ? "success" : "queued"} />
<small>YOLO {datasetValidations[dataset.name]?.pairs.image_label ?? 0} pair · Mask {datasetValidations[dataset.name]?.pairs.image_mask ?? 0} pair</small>
<StatusPill status={datasetValidations[dataset.name]?.ready.mask ? "success" : "queued"} />
<small>Mask {datasetValidations[dataset.name]?.pairs.image_mask ?? 0} pair</small>
</div>
<div className="sampleStrip">
{["images", "labels", "masks"].flatMap((kind) =>
{(["images", "masks"] as const).flatMap((kind) =>
(dataset.samples[kind] ?? []).slice(0, 4).map((sample) => (
<a key={`${kind}-${sample.relative_path}`} href={`${API_BASE}/api/artifacts/${sample.relative_path}`} target="_blank" rel="noreferrer">
<a key={`${kind}-${sample.relative_path}`} href={artifactUrl(sample.relative_path)} target="_blank" rel="noreferrer">
{sample.previewable && <img src={artifactUrl(sample.relative_path)} alt={sample.name} />}
<span>{kind}</span>
<small>{sample.name}</small>
</a>
@@ -1317,29 +1340,6 @@ function App() {
</div>
</section>
<section className="grid two" data-page-section="datasets">
<DatasetPipelineGuide selectedDataset={selectedDataset} validation={selectedValidation} />
<div className="panel">
<div className="panelHead">
<div>
<p className="eyebrow">Current Layout</p>
<h2></h2>
</div>
<FileSearch size={22} />
</div>
{selectedDataset ? (
<div className="pathStack">
<div><span>images</span><code>{selectedDataset.absolute_layout?.images}</code></div>
<div><span>labels</span><code>{selectedDataset.absolute_layout?.labels}</code></div>
<div><span>masks</span><code>{selectedDataset.absolute_layout?.masks}</code></div>
<div><span>pairing</span><code> stem sample.png + sample.txt + sample.png mask</code></div>
</div>
) : (
<p className="muted"></p>
)}
</div>
</section>
<section className="grid two" id="inference" data-page-section="inference">
<div className="panel">
<div className="panelHead">
@@ -1783,45 +1783,6 @@ function JobDiagnostics({ job }: { job: Job }) {
);
}
function DatasetPipelineGuide({ selectedDataset, validation }: { selectedDataset?: UploadedDataset; validation?: DatasetValidation }) {
const activeName = selectedDataset?.name ?? "未选择";
return (
<div className="panel pipelinePanel">
<div className="panelHead">
<div>
<p className="eyebrow">Label Pipeline</p>
<h2> Label </h2>
</div>
<Database size={22} />
</div>
<div className="pipelineExample">
<div>
<span></span>
<strong>DataSet_Own/A_Ori + A_Label_Ori</strong>
<small> A_Label_pro_label_fold A_Label_GT_label_fold</small>
</div>
<div>
<span>YOLO </span>
<strong>Seg_All_In_One_YoloModel/Yolo数据集构建</strong>
<small>ORI/Label ORI_GT_label_fold Data/images Data/labels/*.txt</small>
</div>
</div>
<div className="pipelineSteps">
<div><span>1</span><strong></strong><small> `xxx.png` `xxx.png` label `_label` </small></div>
<div><span>2</span><strong> label pro label</strong><small>`4_deal_labels.py` 使 `*_label.png`</small></div>
<div><span>3</span><strong>pro label mask</strong><small> `Annotate_PALETTE` RGB 0 1 GT mask</small></div>
<div><span>4</span><strong>YOLO txt</strong><small>`2_Check_and_Gen_Txt_Label_sort_label.py` GT polygon </small></div>
</div>
<div className="pipelineStats">
<div><span></span><strong>{activeName}</strong></div>
<div><span>Image/Label</span><strong>{validation?.pairs.image_label ?? 0}</strong></div>
<div><span>Image/Mask</span><strong>{validation?.pairs.image_mask ?? 0}</strong></div>
<div><span>YOLO Ready</span><strong>{validation?.ready.yolo ? "OK" : "Check"}</strong></div>
</div>
</div>
);
}
function WeightPanel({
catalog,
manifest,
@@ -2011,13 +1972,13 @@ function DatasetQuality({ validation }: { validation: DatasetValidation }) {
<div className="qualityBox">
<div className="qualityHead">
<strong>{validation.dataset}</strong>
<span>{validation.ready.yolo ? "YOLO READY" : validation.ready.mask ? "MASK READY" : "CHECK"}</span>
<span>{validation.ready.mask ? "MASK READY" : "CHECK"}</span>
</div>
<div className="qualityStats">
<div><span>Image/Label</span><strong>{validation.pairs.image_label}</strong></div>
<div><span>Image/Mask</span><strong>{validation.pairs.image_mask}</strong></div>
<div><span>Classes</span><strong>{validation.classes.length || 0}</strong></div>
<div><span>Annotations</span><strong>{validation.counts.annotations}</strong></div>
<div><span>Images</span><strong>{validation.counts.images}</strong></div>
</div>
<div className="qualityChecks">
{validation.checks.map((check) => (

View File

@@ -720,7 +720,7 @@ textarea {
.segmented {
display: grid;
grid-template-columns: repeat(3, 1fr);
grid-template-columns: repeat(2, 1fr);
gap: 6px;
padding: 4px;
border: 1px solid var(--line);
@@ -742,7 +742,7 @@ textarea {
}
.drop {
min-height: 118px;
min-height: 128px;
display: grid;
place-items: center;
gap: 8px;
@@ -755,6 +755,14 @@ textarea {
overflow: hidden;
}
.drop small {
max-width: 100%;
padding: 0 12px;
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
}
.drop input {
position: absolute;
inset: 0;
@@ -1073,6 +1081,16 @@ textarea {
background: #0b0d0b;
}
.sampleStrip img {
display: block;
width: 100%;
aspect-ratio: 1.4 / 1;
margin-bottom: 7px;
border-radius: 4px;
object-fit: cover;
background: #060806;
}
.sampleStrip span,
.sampleStrip small {
display: block;