Add dataset QA and custom YOLO training flow

2026-06-30 14:04:11 +08:00
parent 43ed767b4f
commit 93af8bcd3a
14 changed files with 529 additions and 18 deletions
--- a/README.md
+++ b/README.md
@@ -45,9 +45,11 @@ The web UI includes a dataset bench for creating upload workspaces, uploading
 images/labels/masks, and jumping into the existing rename, PNG conversion,
 resize, pair-check, label rebuild, transparent overlay, stitch, and video-frame
 jobs. Selecting an uploaded dataset fills task JSON with its images, labels,
-and masks directories. Segmentation previews, YOLO heatmaps, and loss/metric
-artifacts are grouped on the results dashboard, and YOLO-style `results.csv`
-files are parsed into lightweight training curves.
+and masks directories. The dataset panel validates image/label/mask pairing,
+checks YOLO txt labels and mask dimensions, and can generate a `dataset.yaml`
+for the `yolo.train_custom` task. Segmentation previews, YOLO heatmaps, and
+loss/metric artifacts are grouped on the results dashboard, and YOLO-style
+`results.csv` files are parsed into lightweight training curves.

 The coverage panel calls `GET /api/coverage` and verifies that the user-facing
 scripts from the existing `Seg/` workspace are mapped to web jobs. MMSeg
@@ -119,7 +121,8 @@ The backend exposes all current Seg capabilities as job types. Examples:
  `segmodel.batch_predict`, `segmodel.flops`, `segmodel.params_flops`,
  `segmodel.benchmark`, `segmodel.raw_mask_check`
 - `yolo.train`, `yolo.batch_train`, `yolo.predict`, `yolo.batch_predict`,
-  `yolo.heatmap`, `yolo.compare`, `yolo.raw_mask_check`, `yolo.video_visible`
+  `yolo.train_custom`, `yolo.heatmap`, `yolo.compare`,
+  `yolo.raw_mask_check`, `yolo.video_visible`
 - `mmseg.generate_data`, `mmseg.generate_alg`, `mmseg.train`,
  `mmseg.metrics`, `mmseg.flops_fps`, `mmseg.draw`, `mmseg.extract_loss_miou`
 - `visual.train`, `visual.inference`, `visual.fps`,
--- a/backend/app/acceptance.py
+++ b/backend/app/acceptance.py
@@ -77,11 +77,13 @@ SEGMODEL_TRAIN_STEP_SNIPPET = (


 def _yolo_tiny_train_snippet(root: Path, weight: Path) -> str:
+    custom_train = settings.project_root / "backend" / "app" / "modules" / "yolo" / "custom_train.py"
+    yolo_dir = settings.source_root / "Seg_All_In_One_YoloModel"
    return (
-        "import shutil, cv2, numpy as np; "
+        "import shutil, subprocess, sys, cv2, numpy as np; "
        "from pathlib import Path; "
-        "from ultralytics import YOLO; "
        f"root=Path({str(root)!r}); weight={str(weight)!r}; "
+        f"custom_train=Path({str(custom_train)!r}); yolo_dir=Path({str(yolo_dir)!r}); "
        "shutil.rmtree(root, ignore_errors=True); "
        "[ (root / item).mkdir(parents=True, exist_ok=True) for item in ['images/train','images/val','labels/train','labels/val','runs'] ]; "
        "image=np.zeros((64,64,3), dtype=np.uint8); "
@@ -91,8 +93,7 @@ def _yolo_tiny_train_snippet(root: Path, weight: Path) -> str:
        "    cv2.imwrite(str(root / 'images' / split / 'sample.jpg'), image)\n"
        "    (root / 'labels' / split / 'sample.txt').write_text(label, encoding='utf-8')\n"
        "(root / 'data.yaml').write_text('path: '+str(root)+'\\ntrain: images/train\\nval: images/val\\nnc: 1\\nnames:\\n  0: object\\n', encoding='utf-8'); "
-        "model=YOLO(weight); "
-        "model.train(data=str(root/'data.yaml'), epochs=1, imgsz=64, batch=1, workers=0, device='cpu', project=str(root/'runs'), name='tiny', exist_ok=True, verbose=False, plots=False, val=False); "
+        "subprocess.run([sys.executable, str(custom_train), '--data', str(root/'data.yaml'), '--model', str(weight), '--epochs', '1', '--imgsz', '64', '--batch', '1', '--workers', '0', '--device', 'cpu', '--project', str(root/'runs'), '--name', 'tiny', '--exist-ok'], cwd=str(yolo_dir), check=True); "
        "results=root/'runs'/'tiny'/'results.csv'; best=root/'runs'/'tiny'/'weights'/'best.pt'; "
        "assert results.exists() and results.stat().st_size > 0; "
        "assert best.exists() and best.stat().st_size > 0; "
@@ -418,16 +419,46 @@ def run_live_acceptance(base_url: str = "http://127.0.0.1:8010") -> dict[str, An
    created_dataset = _request_json("POST", f"{base_url}/api/datasets", {"name": dataset_name, "description": "acceptance smoke"}, timeout=10)
    checks.append({"name": "create_dataset_api", "passed": created_dataset.get("passed", False), "detail": created_dataset})

+    import cv2
+    import numpy as np
+
+    image = np.zeros((16, 16, 3), dtype=np.uint8)
+    image[:, :, 1] = 160
+    mask = np.zeros((16, 16), dtype=np.uint8)
+    mask[4:12, 4:12] = 255
+    _, image_encoded = cv2.imencode(".png", image)
+    _, mask_encoded = cv2.imencode(".png", mask)
+
+    upload_image = _post_multipart(
+        f"{base_url}/api/datasets/{dataset_name}/upload/images",
+        "files",
+        "sample.png",
+        image_encoded.tobytes(),
+        "image/png",
+        timeout=10,
+    )
+    checks.append({"name": "upload_image_api", "passed": upload_image.get("passed", False), "detail": upload_image})
+
    upload = _post_multipart(
        f"{base_url}/api/datasets/{dataset_name}/upload/labels",
        "files",
-        "label 01.txt",
+        "sample.txt",
        b"0 0.5 0.5 0.25 0.25\n",
        "text/plain",
        timeout=10,
    )
    checks.append({"name": "upload_label_api", "passed": upload.get("passed", False), "detail": upload})

+    upload_mask = _post_multipart(
+        f"{base_url}/api/datasets/{dataset_name}/upload/masks",
+        "files",
+        "sample.png",
+        mask_encoded.tobytes(),
+        "image/png",
+        timeout=10,
+    )
+    checks.append({"name": "upload_mask_api", "passed": upload_mask.get("passed", False), "detail": upload_mask})
+
    artifact_ok = False
    artifact_detail: dict[str, Any] = {"skipped": True}
    try:
@@ -438,6 +469,26 @@ def run_live_acceptance(base_url: str = "http://127.0.0.1:8010") -> dict[str, An
        artifact_detail = {"error": str(exc)}
    checks.append({"name": "artifact_api_for_uploaded_label", "passed": artifact_ok, "detail": artifact_detail})

+    dataset_validation = _request_json("GET", f"{base_url}/api/datasets/{dataset_name}/validate", timeout=10)
+    validation_json = dataset_validation.get("json") if dataset_validation.get("passed") else {}
+    checks.append(
+        {
+            "name": "dataset_validate_api",
+            "passed": dataset_validation.get("passed", False) and validation_json.get("ready", {}).get("yolo") and validation_json.get("ready", {}).get("mask"),
+            "detail": dataset_validation,
+        }
+    )
+
+    yolo_yaml = _request_json("POST", f"{base_url}/api/datasets/{dataset_name}/yolo-yaml", {"class_names": ["object"]}, timeout=10)
+    yolo_yaml_json = yolo_yaml.get("json") if yolo_yaml.get("passed") else {}
+    checks.append(
+        {
+            "name": "dataset_yolo_yaml_api",
+            "passed": yolo_yaml.get("passed", False) and "dataset.yaml" in str(yolo_yaml_json.get("relative_path", "")),
+            "detail": yolo_yaml,
+        }
+    )
+
    mock = _create_job_and_wait(base_url, "mock.echo", {"message": f"acceptance {run_id}"}, timeout=45)
    mock_log = mock.get("polled", {}).get("job", {}).get("log_tail", "")
    checks.append({"name": "mock_job_runner", "passed": mock.get("passed", False) and f"acceptance {run_id}" in mock_log, "detail": mock})
--- a/backend/app/agents/evaluation_agent.py
+++ b/backend/app/agents/evaluation_agent.py
@@ -40,14 +40,17 @@ def evaluate_project() -> dict:
    expectations = {
        "left_nav_dataset": "数据集" in frontend_text and "#datasets" in frontend_text,
        "upload_ui": "uploadDatasetFiles" in frontend_text and "labels" in frontend_text and "masks" in frontend_text,
+        "dataset_quality_ui": "DatasetQuality" in frontend_text and "generateSelectedYoloYaml" in frontend_text,
        "loss_result_ui": "loss" in frontend_text.lower() and "heatmap" in frontend_text.lower() and "CurvePanel" in frontend_text,
        "dataset_api": "/api/datasets" in backend_text and "api_upload_dataset_files" in backend_text,
+        "dataset_quality_api": "/api/datasets/{dataset_name}/validate" in backend_text and "/api/datasets/{dataset_name}/yolo-yaml" in backend_text,
        "curve_api": "/api/results/curves" in backend_text,
        "deep_acceptance_api": "/api/acceptance/deep" in backend_text,
        "deep_acceptance_ui": "runDeepAcceptance" in frontend_text and "深度训练" in frontend_text,
        "deep_yolo_heatmap_validation": "yolo_tiny_heatmap_generation" in acceptance_text,
        "coverage_api": "/api/coverage" in backend_text and coverage["task_build_passed"],
        "visual_tools": "visual.yolo11_heatmap_v2" in catalog["task_types"] and "visual.fps" in catalog["task_types"],
+        "yolo_custom_train": "yolo.train_custom" in catalog["task_types"],
        "yolo_dataset_tools": "dataset.yolo_txt_sort" in catalog["task_types"] and "dataset.yolo_resize" in catalog["task_types"],
        "no_weight_to_gitea": "Do not push" in readme_text and "check_no_weight_git" in readme_text,
        "all_core_tasks": all(task in catalog["task_types"] for task in REQUIRED_TASKS if REQUIRED_TASKS[task] == "job"),
--- a/backend/app/agents/validation_agent.py
+++ b/backend/app/agents/validation_agent.py
@@ -48,6 +48,7 @@ def validate_project(run_build: bool = False) -> dict:
    coverage = get_coverage_report()

    checks.append({"name": "catalog_has_yolo_heatmap", "passed": "yolo.heatmap" in catalog["task_types"]})
+    checks.append({"name": "catalog_has_yolo_custom_train", "passed": "yolo.train_custom" in catalog["task_types"]})
    checks.append({"name": "catalog_has_mmseg_31_algs", "passed": len(catalog["mmseg_algorithms"]) >= 31})
    checks.append({"name": "catalog_has_visual_tools", "passed": "visual.yolo11_heatmap_v2" in catalog["task_types"] and "visual.fps" in catalog["task_types"]})
    checks.append({"name": "catalog_has_yolo_dataset_tools", "passed": "dataset.yolo_txt_sort" in catalog["task_types"] and "dataset.yolo_convert_png" in catalog["task_types"]})
--- a/backend/app/catalog.py
+++ b/backend/app/catalog.py
@@ -81,6 +81,7 @@ TASK_TYPES = [
    "segmodel.metrics",
    "segmodel.copy_best",
    "yolo.train",
+    "yolo.train_custom",
    "yolo.batch_train",
    "yolo.predict",
    "yolo.predict_v1",
@@ -181,6 +182,7 @@ TASK_DEFAULTS: dict[str, dict[str, Any]] = {
    "segmodel.metrics": {},
    "segmodel.copy_best": {},
    "yolo.train": {"model": "YOLOv8n-seg"},
+    "yolo.train_custom": {"data": "var/uploads/datasets/example/dataset.yaml", "model": "YOLO11n-seg", "epochs": 10, "imgsz": 640, "batch": 1, "workers": 0, "device": "cpu", "exist_ok": True},
    "yolo.batch_train": {},
    "yolo.predict": {"model": "YOLOv8n-seg", "pt_name": "best.pt", "conf": 0.2, "run_choice": 1},
    "yolo.predict_v1": {"model": "YOLOv8n-seg", "pt_name": "best.pt", "conf": 0.2, "run_choice": 1},
--- a/backend/app/main.py
+++ b/backend/app/main.py
@@ -16,12 +16,12 @@ from .coverage import get_coverage_report
 from .jobs import cancel_job, create_job
 from .modules.results.service import scan_results, scan_training_curves
 from .modules.system.service import disk_usage, get_conda_envs, get_gpus
-from .modules.dataset.service import create_dataset, list_uploaded_datasets, save_upload
+from .modules.dataset.service import create_dataset, generate_yolo_dataset_yaml, list_uploaded_datasets, save_upload, validate_dataset
 from .modules.weights.service import load_manifest, sync_weights, verify_weights
 from .agents.evaluation_agent import evaluate_project
 from .agents.validation_agent import validate_project
 from .paths import ensure_inside
-from .schemas import DatasetCreate, JobCreate, ProfileCreate, WeightSyncRequest
+from .schemas import DatasetCreate, DatasetYoloYamlRequest, JobCreate, ProfileCreate, WeightSyncRequest

 app = FastAPI(title="Seg Data Server Net", version="0.1.0")

@@ -109,6 +109,19 @@ async def api_upload_dataset_files(dataset_name: str, kind: str, files: list[Upl
        raise HTTPException(status_code=400, detail=str(exc)) from exc


+@app.get("/api/datasets/{dataset_name}/validate")
+def api_validate_dataset(dataset_name: str) -> dict:
+    return validate_dataset(dataset_name)
+
+
+@app.post("/api/datasets/{dataset_name}/yolo-yaml")
+def api_generate_dataset_yolo_yaml(dataset_name: str, request: DatasetYoloYamlRequest) -> dict:
+    try:
+        return generate_yolo_dataset_yaml(dataset_name, request.class_names)
+    except Exception as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+
+
@app.get("/api/profiles")
 def api_profiles(kind: str | None = None) -> list[dict]:
    return db.list_profiles(kind)
--- a/backend/app/modules/dataset/service.py
+++ b/backend/app/modules/dataset/service.py
@@ -13,6 +13,7 @@ from ...config import settings

 DATASET_KINDS = ("images", "labels", "masks")
 IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}
+LABEL_EXTS = {".txt", ".json", ".yaml", ".yml"}


 def uploads_root() -> Path:
@@ -76,6 +77,62 @@ def _iter_files(root: Path) -> Iterable[Path]:
    return sorted(path for path in root.rglob("*") if path.is_file())


+def _stem_map(paths: Iterable[Path], exts: set[str] | None = None) -> dict[str, Path]:
+    result: dict[str, Path] = {}
+    for path in paths:
+        if exts and path.suffix.lower() not in exts:
+            continue
+        result.setdefault(path.stem, path)
+    return result
+
+
+def _image_shape(path: Path) -> dict | None:
+    try:
+        import cv2
+
+        image = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
+        if image is None:
+            return None
+        height, width = image.shape[:2]
+        channels = 1 if image.ndim == 2 else image.shape[2]
+        return {"width": int(width), "height": int(height), "channels": int(channels)}
+    except Exception:
+        return None
+
+
+def _validate_yolo_txt(path: Path) -> dict:
+    errors = []
+    classes: set[int] = set()
+    lines = path.read_text(encoding="utf-8", errors="replace").splitlines()
+    annotation_count = 0
+    for line_number, line in enumerate(lines, 1):
+        raw = line.strip()
+        if not raw:
+            continue
+        parts = raw.split()
+        annotation_count += 1
+        if len(parts) < 5:
+            errors.append(f"{path.name}:{line_number} has fewer than 5 tokens")
+            continue
+        try:
+            class_id = int(float(parts[0]))
+            coords = [float(item) for item in parts[1:]]
+        except ValueError:
+            errors.append(f"{path.name}:{line_number} contains non-numeric values")
+            continue
+        if class_id < 0:
+            errors.append(f"{path.name}:{line_number} has negative class id")
+        classes.add(class_id)
+        if len(coords) not in {4} and len(coords) < 6:
+            errors.append(f"{path.name}:{line_number} has too few segmentation coordinates")
+        if len(coords) != 4 and len(coords) % 2 != 0:
+            errors.append(f"{path.name}:{line_number} has an odd number of polygon coordinates")
+        out_of_range = [value for value in coords if value < 0 or value > 1]
+        if out_of_range:
+            errors.append(f"{path.name}:{line_number} has coordinates outside 0..1")
+    return {"annotations": annotation_count, "classes": sorted(classes), "errors": errors}
+
+
 def describe_dataset(name: str) -> dict:
    safe_name = slugify(name)
    root = dataset_dir(safe_name)
@@ -99,6 +156,113 @@ def describe_dataset(name: str) -> dict:
    return {**meta, "absolute_layout": absolute_layout, "counts": counts, "samples": samples}


+def validate_dataset(name: str) -> dict:
+    safe_name = slugify(name)
+    root = dataset_dir(safe_name)
+    image_files = list(_iter_files(root / "images"))
+    label_files = list(_iter_files(root / "labels"))
+    mask_files = list(_iter_files(root / "masks"))
+    images = _stem_map(image_files, IMAGE_EXTS)
+    labels = _stem_map(label_files, LABEL_EXTS)
+    masks = _stem_map(mask_files, IMAGE_EXTS)
+
+    image_stems = set(images)
+    label_stems = set(labels)
+    mask_stems = set(masks)
+    paired_label_stems = sorted(image_stems & label_stems)
+    paired_mask_stems = sorted(image_stems & mask_stems)
+
+    yolo_errors = []
+    class_ids: set[int] = set()
+    annotation_count = 0
+    for label in labels.values():
+        if label.suffix.lower() != ".txt":
+            continue
+        detail = _validate_yolo_txt(label)
+        yolo_errors.extend(detail["errors"])
+        class_ids.update(detail["classes"])
+        annotation_count += detail["annotations"]
+
+    shape_mismatches = []
+    sample_shapes = []
+    for stem in paired_mask_stems[:80]:
+        image_shape = _image_shape(images[stem])
+        mask_shape = _image_shape(masks[stem])
+        if image_shape and len(sample_shapes) < 8:
+            sample_shapes.append({"name": images[stem].name, **image_shape})
+        if image_shape and mask_shape and (image_shape["width"], image_shape["height"]) != (mask_shape["width"], mask_shape["height"]):
+            shape_mismatches.append({"stem": stem, "image": image_shape, "mask": mask_shape})
+    if not sample_shapes:
+        for image in list(images.values())[:8]:
+            image_shape = _image_shape(image)
+            if image_shape:
+                sample_shapes.append({"name": image.name, **image_shape})
+
+    checks = [
+        {"name": "has_images", "passed": len(images) > 0, "count": len(images)},
+        {"name": "has_labels_or_masks", "passed": len(labels) > 0 or len(masks) > 0, "labels": len(labels), "masks": len(masks)},
+        {"name": "image_label_pairs", "passed": len(label_stems) == 0 or len(paired_label_stems) > 0, "count": len(paired_label_stems)},
+        {"name": "image_mask_pairs", "passed": len(mask_stems) == 0 or len(paired_mask_stems) > 0, "count": len(paired_mask_stems)},
+        {"name": "yolo_txt_valid", "passed": not yolo_errors, "errors": yolo_errors[:20]},
+        {"name": "mask_shapes_match", "passed": not shape_mismatches, "errors": shape_mismatches[:20]},
+    ]
+    yolo_ready = len(images) > 0 and len(paired_label_stems) > 0 and not yolo_errors
+    mask_ready = len(images) > 0 and len(paired_mask_stems) > 0 and not shape_mismatches
+    return {
+        "dataset": safe_name,
+        "root": str(root.resolve()),
+        "counts": {"images": len(images), "labels": len(labels), "masks": len(masks), "annotations": annotation_count},
+        "pairs": {
+            "image_label": len(paired_label_stems),
+            "image_mask": len(paired_mask_stems),
+            "images_without_labels": sorted(image_stems - label_stems)[:50],
+            "labels_without_images": sorted(label_stems - image_stems)[:50],
+            "images_without_masks": sorted(image_stems - mask_stems)[:50],
+            "masks_without_images": sorted(mask_stems - image_stems)[:50],
+        },
+        "classes": sorted(class_ids),
+        "sample_shapes": sample_shapes,
+        "checks": checks,
+        "ready": {"yolo": yolo_ready, "mask": mask_ready, "any": yolo_ready or mask_ready},
+    }
+
+
+def generate_yolo_dataset_yaml(name: str, class_names: list[str] | None = None) -> dict:
+    validation = validate_dataset(name)
+    if not validation["ready"]["yolo"]:
+        raise ValueError("dataset is not YOLO-ready; upload matching images and valid txt labels first")
+    safe_name = slugify(name)
+    root = dataset_dir(safe_name)
+    classes = validation["classes"] or [0]
+    class_count = max(classes) + 1
+    names = list(class_names or [])
+    if len(names) < class_count:
+        names.extend(f"class_{index}" for index in range(len(names), class_count))
+    yaml_path = root / "dataset.yaml"
+    names_block = "\n".join(f"  {index}: {label}" for index, label in enumerate(names))
+    yaml_text = "\n".join(
+        [
+            f"path: {root.resolve()}",
+            "train: images",
+            "val: images",
+            f"nc: {len(names)}",
+            "names:",
+            names_block,
+            "",
+        ]
+    )
+    yaml_path.write_text(yaml_text, encoding="utf-8")
+    return {
+        "dataset": safe_name,
+        "path": str(yaml_path.resolve()),
+        "relative_path": str(yaml_path.resolve().relative_to(settings.project_root)),
+        "classes": classes,
+        "names": names,
+        "content": yaml_text,
+        "validation": validation,
+    }
+
+
 def list_uploaded_datasets() -> list[dict]:
    root = uploads_root()
    datasets = []
--- a/backend/app/modules/yolo/custom_train.py
+++ b/backend/app/modules/yolo/custom_train.py
@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from ultralytics import YOLO
+
+
+MODEL_ALIASES = {
+    "YOLOv8n-seg": "yolov8n-seg.pt",
+    "YOLOv8s-seg": "yolov8s-seg.pt",
+    "YOLOv8m-seg": "yolov8m-seg.pt",
+    "YOLOv8l-seg": "yolov8l-seg.pt",
+    "YOLOv8x-seg": "yolov8x-seg.pt",
+    "YOLOv9c-seg": "yolov9c-seg.pt",
+    "YOLOv9e-seg": "yolov9e-seg.pt",
+    "YOLO11n-seg": "yolo11n-seg.pt",
+    "YOLO11s-seg": "yolo11s-seg.pt",
+    "YOLO11m-seg": "yolo11m-seg.pt",
+    "YOLO11l-seg": "yolo11l-seg.pt",
+    "YOLO11x-seg": "yolo11x-seg.pt",
+}
+
+
+def resolve_model(value: str) -> str:
+    candidate = MODEL_ALIASES.get(value, value)
+    path = Path(candidate).expanduser()
+    if path.is_absolute() or path.exists():
+        return str(path.resolve())
+    return candidate
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Train a YOLO segmentation model from a supplied dataset.yaml.")
+    parser.add_argument("--data", required=True, help="Path to YOLO dataset.yaml.")
+    parser.add_argument("--model", default="YOLO11n-seg", help="Model alias, weight path, or Ultralytics model name.")
+    parser.add_argument("--epochs", type=int, default=10)
+    parser.add_argument("--imgsz", type=int, default=640)
+    parser.add_argument("--batch", type=int, default=1)
+    parser.add_argument("--workers", type=int, default=0)
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--project", required=True)
+    parser.add_argument("--name", default="custom_upload")
+    parser.add_argument("--exist-ok", action="store_true")
+    args = parser.parse_args()
+
+    model = YOLO(resolve_model(args.model))
+    result = model.train(
+        data=str(Path(args.data).expanduser().resolve()),
+        epochs=args.epochs,
+        imgsz=args.imgsz,
+        batch=args.batch,
+        workers=args.workers,
+        device=args.device,
+        project=str(Path(args.project).expanduser().resolve()),
+        name=args.name,
+        exist_ok=args.exist_ok,
+        verbose=True,
+    )
+    save_dir = getattr(result, "save_dir", None) or getattr(model.trainer, "save_dir", "")
+    print(f"save_dir={save_dir}")
+
+
+if __name__ == "__main__":
+    main()
--- a/backend/app/modules/yolo/tasks.py
+++ b/backend/app/modules/yolo/tasks.py
@@ -6,6 +6,7 @@ from ...config import settings

 YOLO_DIR = settings.source_root / "Seg_All_In_One_YoloModel"
 VIDEO_YOLO_DIR = settings.source_root / "Seg_Predict_YoloModel"
+CUSTOM_TRAIN = settings.project_root / "backend" / "app" / "modules" / "yolo" / "custom_train.py"


 def build_yolo_task(job_type: str, params: dict, conda_env: str) -> CommandSpec | None:
@@ -16,6 +17,21 @@ def build_yolo_task(job_type: str, params: dict, conda_env: str) -> CommandSpec
        append_flag(args, "--model", required(params, "model"))
        return CommandSpec(args, YOLO_DIR, "train one Ultralytics YOLO segmentation model")

+    if job_type == "yolo.train_custom":
+        args = conda_python(conda_env, CUSTOM_TRAIN)
+        append_flag(args, "--data", required(params, "data"))
+        append_flag(args, "--model", params.get("model", "YOLO11n-seg"))
+        append_flag(args, "--epochs", params.get("epochs", 10))
+        append_flag(args, "--imgsz", params.get("imgsz", 640))
+        append_flag(args, "--batch", params.get("batch", 1))
+        append_flag(args, "--workers", params.get("workers", 0))
+        append_flag(args, "--device", params.get("device", "cpu"))
+        append_flag(args, "--project", params.get("project", settings.project_root / "var" / "custom_yolo_runs"))
+        append_flag(args, "--name", params.get("name", "custom_upload"))
+        if params.get("exist_ok", True):
+            args.append("--exist-ok")
+        return CommandSpec(args, YOLO_DIR, "train YOLO segmentation on a supplied dataset.yaml")
+
    if job_type == "yolo.batch_train":
        return CommandSpec(bash(YOLO_DIR / "yolo_train.sh"), YOLO_DIR, "run legacy YOLO batch training", env=env)

--- a/backend/app/schemas.py
+++ b/backend/app/schemas.py
@@ -47,3 +47,7 @@ class WeightSyncRequest(BaseModel):
 class DatasetCreate(BaseModel):
    name: str
    description: str = ""
+
+
+class DatasetYoloYamlRequest(BaseModel):
+    class_names: list[str] | None = None
--- a/backend/tests/test_catalog.py
+++ b/backend/tests/test_catalog.py
@@ -7,9 +7,9 @@ def test_catalog_contains_required_capabilities():
    for task in [
        "dataset.video_frames",
        "segmodel.train",
+        "yolo.train_custom",
        "yolo.predict",
        "mmseg.flops_fps",
        "analysis.all",
    ]:
        assert task in TASK_TYPES
-
--- a/backend/tests/test_dataset_service.py
+++ b/backend/tests/test_dataset_service.py
@@ -1,4 +1,7 @@
-from app.modules.dataset.service import create_dataset, describe_dataset
+import cv2
+import numpy as np
+
+from app.modules.dataset.service import create_dataset, describe_dataset, generate_yolo_dataset_yaml, validate_dataset


 def test_create_dataset_layout(tmp_path, monkeypatch):
@@ -11,3 +14,29 @@ def test_create_dataset_layout(tmp_path, monkeypatch):
    assert created["counts"] == {"images": 0, "labels": 0, "masks": 0}
    described = describe_dataset("case_01")
    assert described["layout"]["images"].endswith("images")
+
+
+def test_validate_dataset_and_generate_yolo_yaml(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 yolo", "demo")
+    root = service.dataset_dir("case_yolo")
+    image = np.zeros((32, 32, 3), dtype=np.uint8)
+    mask = np.zeros((32, 32), dtype=np.uint8)
+    cv2.imwrite(str(root / "images" / "sample.png"), image)
+    cv2.imwrite(str(root / "masks" / "sample.png"), mask)
+    (root / "labels" / "sample.txt").write_text("0 0.25 0.25 0.75 0.25 0.75 0.75 0.25 0.75\n", encoding="utf-8")
+
+    validation = validate_dataset("case_yolo")
+
+    assert validation["ready"] == {"yolo": True, "mask": True, "any": True}
+    assert validation["pairs"]["image_label"] == 1
+    assert validation["classes"] == [0]
+
+    generated = generate_yolo_dataset_yaml("case_yolo", ["tool"])
+
+    assert generated["relative_path"] == "var/uploads/datasets/case_yolo/dataset.yaml"
+    assert "nc: 1" in generated["content"]
+    assert "0: tool" in generated["content"]
--- a/frontend/src/main.tsx
+++ b/frontend/src/main.tsx
@@ -56,6 +56,22 @@ type UploadedDataset = {
  samples: Record<string, Array<{ name: string; relative_path: string; size: number; previewable: boolean }>>;
 };

+type DatasetValidation = {
+  dataset: string;
+  counts: { images: number; labels: number; masks: number; annotations: number };
+  pairs: {
+    image_label: number;
+    image_mask: number;
+    images_without_labels: string[];
+    labels_without_images: string[];
+    images_without_masks: string[];
+    masks_without_images: string[];
+  };
+  classes: number[];
+  checks: Array<{ name: string; passed: boolean; count?: number; labels?: number; masks?: number; errors?: unknown[] }>;
+  ready: { yolo: boolean; mask: boolean; any: boolean };
+};
+
 type ResultItem = {
  name: string;
  path: string;
@@ -152,6 +168,7 @@ const defaultParams: Record<string, Record<string, unknown>> = {
  "segmodel.train": { architecture: "Unet" },
  "segmodel.predict": { architecture: "Unet", run_choice: 1 },
  "yolo.train": { model: "YOLOv8n-seg" },
+  "yolo.train_custom": { model: "YOLO11n-seg", data: "var/uploads/datasets/example/dataset.yaml", epochs: 10, imgsz: 640, batch: 1, workers: 0, device: "cpu", exist_ok: true },
  "yolo.predict": { model: "YOLOv8n-seg", pt_name: "best.pt", conf: 0.2, run_choice: 1 },
  "yolo.heatmap": { model: "YOLOv8n-seg", cam_method: "All", pt_name: "best.pt", run_choice: 1 },
  "mmseg.generate_alg": { dataset_choice: 1, gpu_count: 1, gpu_ids: [0], schedule_mode: 2, max_epochs: 300, algorithm_choice: 1 },
@@ -197,6 +214,7 @@ function useData() {
  const [results, setResults] = useState<ResultItem[]>([]);
  const [curves, setCurves] = useState<TrainingCurve[]>([]);
  const [datasets, setDatasets] = useState<UploadedDataset[]>([]);
+  const [datasetValidations, setDatasetValidations] = useState<Record<string, DatasetValidation>>({});
  const [coverage, setCoverage] = useState<CoveragePayload | null>(null);
  const [acceptance, setAcceptance] = useState<AcceptancePayload | null>(null);
  const [deepAcceptance, setDeepAcceptance] = useState<DeepAcceptancePayload | null>(null);
@@ -221,6 +239,18 @@ function useData() {
      setResults(resultsNext.slice(0, 80));
      setCurves(curvesNext.slice(0, 12));
      setDatasets(datasetsNext);
+      const validationEntries: Array<[string, DatasetValidation]> = [];
+      await Promise.all(
+        datasetsNext.map(async (dataset) => {
+          try {
+            const validation = await api<DatasetValidation>(`/api/datasets/${encodeURIComponent(dataset.name)}/validate`);
+            validationEntries.push([dataset.name, validation]);
+          } catch {
+            // Dataset validation is advisory; upload and job controls should remain usable.
+          }
+        })
+      );
+      setDatasetValidations(Object.fromEntries(validationEntries));
      setCoverage(coverageNext);
      setAcceptance(acceptanceNext);
      setDeepAcceptance(deepAcceptanceNext);
@@ -236,7 +266,7 @@ function useData() {
    return () => window.clearInterval(timer);
  }, []);

-  return { catalog, gpus, jobs, results, curves, datasets, coverage, acceptance, deepAcceptance, error, refresh };
+  return { catalog, gpus, jobs, results, curves, datasets, datasetValidations, coverage, acceptance, deepAcceptance, error, refresh };
 }

 function StatusPill({ status }: { status: string }) {
@@ -244,7 +274,7 @@ function StatusPill({ status }: { status: string }) {
 }

 function App() {
-  const { catalog, gpus, jobs, results, curves, datasets, coverage, acceptance, deepAcceptance, error, refresh } = useData();
+  const { catalog, gpus, jobs, results, curves, datasets, datasetValidations, coverage, acceptance, deepAcceptance, error, refresh } = useData();
  const [taskType, setTaskType] = useState("mock.echo");
  const [params, setParams] = useState(JSON.stringify(defaultParams["mock.echo"], null, 2));
  const [selectedJob, setSelectedJob] = useState<Job | null>(null);
@@ -278,6 +308,7 @@ function App() {
    () => datasets.find((dataset) => dataset.name === selectedDatasetName) ?? datasets.find((dataset) => dataset.name === datasetName),
    [datasetName, datasets, selectedDatasetName]
  );
+  const selectedValidation = selectedDataset ? datasetValidations[selectedDataset.name] : undefined;
  const selectedCurve = curves.find((curve) => curve.relative_path === selectedCurvePath) ?? curves[0];

  function pickTask(next: string) {
@@ -385,6 +416,23 @@ function App() {
    }
  }

+  async function generateSelectedYoloYaml() {
+    if (!selectedDataset) return;
+    setBusy(true);
+    try {
+      const classNames = selectedValidation?.classes.map((classId) => `class_${classId}`) ?? undefined;
+      const generated = await api<{ relative_path: string; path: string }>(`/api/datasets/${encodeURIComponent(selectedDataset.name)}/yolo-yaml`, {
+        method: "POST",
+        body: JSON.stringify({ class_names: classNames })
+      });
+      setTaskType("yolo.train_custom");
+      setParams(JSON.stringify({ model: "YOLO11n-seg", data: generated.path, epochs: 10, imgsz: 640, batch: 1, workers: 0, device: "cpu", exist_ok: true }, null, 2));
+      await refresh();
+    } finally {
+      setBusy(false);
+    }
+  }
+
  async function inspectJob(job: Job) {
    const detail = await api<Job>(`/api/jobs/${job.id}`);
    setSelectedJob(detail);
@@ -565,13 +613,18 @@ function App() {
                <p className="eyebrow">Files</p>
                <h2>数据集浏览</h2>
              </div>
-              <FileImage size={22} />
+              <div className="buttonRow compactButtons">
+                <button className="iconButton" disabled={busy || !selectedValidation?.ready.yolo} onClick={generateSelectedYoloYaml} title="生成 YOLO dataset.yaml">
+                  <FileSearch size={18} />
+                </button>
+                <FileImage size={22} />
+              </div>
            </div>
            <div className="datasetList">
              {datasets.map((dataset) => (
+                <div key={dataset.name}>
                  <div
                    className={`datasetCard ${selectedDataset?.name === dataset.name ? "selected" : ""}`}
-                    key={dataset.name}
                    role="button"
                    tabIndex={0}
                    onClick={() => {
@@ -589,6 +642,10 @@ function App() {
                    <strong>{dataset.name}</strong>
                    <span>{dataset.counts.images} image · {dataset.counts.labels} label · {dataset.counts.masks} mask</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>
+                  </div>
                  <div className="sampleStrip">
                    {["images", "labels", "masks"].flatMap((kind) =>
                      (dataset.samples[kind] ?? []).slice(0, 4).map((sample) => (
@@ -600,8 +657,10 @@ function App() {
                    )}
                  </div>
                  </div>
+                </div>
              ))}
            </div>
+            {selectedValidation && <DatasetQuality validation={selectedValidation} />}
          </div>
        </section>

@@ -819,6 +878,31 @@ function ResultPreview({ results }: { results: ResultItem[] }) {
  );
 }

+function DatasetQuality({ validation }: { validation: DatasetValidation }) {
+  return (
+    <div className="qualityBox">
+      <div className="qualityHead">
+        <strong>{validation.dataset}</strong>
+        <span>{validation.ready.yolo ? "YOLO READY" : 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>
+      <div className="qualityChecks">
+        {validation.checks.map((check) => (
+          <div key={check.name} className={check.passed ? "ok" : "bad"}>
+            <span>{check.name}</span>
+            <small>{check.passed ? "ok" : `${check.errors?.length ?? 0} issue`}</small>
+          </div>
+        ))}
+      </div>
+    </div>
+  );
+}
+
 function CurvePanel({
  curves,
  selected,
--- a/frontend/src/styles.css
+++ b/frontend/src/styles.css
@@ -538,6 +538,79 @@ textarea {
  font-size: 11px;
 }

+.readinessLine {
+  display: flex;
+  align-items: center;
+  gap: 8px;
+  margin-bottom: 10px;
+}
+
+.qualityBox {
+  display: grid;
+  gap: 10px;
+  margin-top: 12px;
+  padding: 12px;
+  border: 1px solid var(--line);
+  border-radius: 7px;
+  background: #0b0d0b;
+}
+
+.qualityHead {
+  display: flex;
+  justify-content: space-between;
+  gap: 12px;
+}
+
+.qualityHead span {
+  color: var(--green);
+  font-size: 12px;
+  font-weight: 760;
+}
+
+.qualityStats {
+  display: grid;
+  grid-template-columns: repeat(4, minmax(0, 1fr));
+  gap: 8px;
+}
+
+.qualityStats div,
+.qualityChecks div {
+  min-width: 0;
+  padding: 8px;
+  border: 1px solid var(--line);
+  border-radius: 6px;
+  background: #101310;
+}
+
+.qualityStats span,
+.qualityStats strong,
+.qualityChecks span,
+.qualityChecks small {
+  display: block;
+  overflow: hidden;
+  text-overflow: ellipsis;
+  white-space: nowrap;
+}
+
+.qualityStats span {
+  color: var(--muted);
+  font-size: 11px;
+}
+
+.qualityChecks {
+  display: grid;
+  grid-template-columns: repeat(2, minmax(0, 1fr));
+  gap: 8px;
+}
+
+.qualityChecks .ok {
+  border-color: rgba(157, 226, 111, 0.32);
+}
+
+.qualityChecks .bad {
+  border-color: rgba(240, 113, 103, 0.55);
+}
+
 .jobList, .resultList {
  display: grid;
  gap: 8px;
@@ -573,6 +646,7 @@ textarea {
 .pill-success { color: var(--green); }
 .pill-failed { color: var(--red); }
 .pill-cancelled { color: var(--amber); }
+.pill-queued { color: var(--muted); }

 .gpu {
  display: grid;
@@ -758,7 +832,9 @@ meter {

  .opGrid,
  .sampleStrip,
-  .taskCheckList {
+  .taskCheckList,
+  .qualityStats,
+  .qualityChecks {
    grid-template-columns: repeat(2, minmax(0, 1fr));
  }
 }