Add dataset QA and custom YOLO training flow

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})

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"),

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"]})

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},

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)

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 = []

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()

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)
 

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

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
-

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"]