diff --git a/README.md b/README.md index 579cb71..8381819 100644 --- 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`, diff --git a/backend/app/acceptance.py b/backend/app/acceptance.py index dd212de..7f16545 100644 --- 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}) diff --git a/backend/app/agents/evaluation_agent.py b/backend/app/agents/evaluation_agent.py index 1721cd4..345b23a 100644 --- 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"), diff --git a/backend/app/agents/validation_agent.py b/backend/app/agents/validation_agent.py index ef849ed..62c972f 100644 --- 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"]}) diff --git a/backend/app/catalog.py b/backend/app/catalog.py index 99be1de..9610950 100644 --- 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}, diff --git a/backend/app/main.py b/backend/app/main.py index 7198599..c26fd1a 100644 --- 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) diff --git a/backend/app/modules/dataset/service.py b/backend/app/modules/dataset/service.py index 93b0a8f..136d32a 100644 --- 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 = [] diff --git a/backend/app/modules/yolo/custom_train.py b/backend/app/modules/yolo/custom_train.py new file mode 100644 index 0000000..f6b304e --- /dev/null +++ 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() diff --git a/backend/app/modules/yolo/tasks.py b/backend/app/modules/yolo/tasks.py index 972bc96..f84d4ad 100644 --- 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) diff --git a/backend/app/schemas.py b/backend/app/schemas.py index b1cbfdc..cd62c1f 100644 --- 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 diff --git a/backend/tests/test_catalog.py b/backend/tests/test_catalog.py index 2b3967c..4e2afd3 100644 --- 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 - diff --git a/backend/tests/test_dataset_service.py b/backend/tests/test_dataset_service.py index 88f7435..20272cb 100644 --- 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"] diff --git a/frontend/src/main.tsx b/frontend/src/main.tsx index 09e8cf3..bcdf601 100644 --- a/frontend/src/main.tsx +++ b/frontend/src/main.tsx @@ -56,6 +56,22 @@ type UploadedDataset = { samples: Record>; }; +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> = { "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([]); const [curves, setCurves] = useState([]); const [datasets, setDatasets] = useState([]); + const [datasetValidations, setDatasetValidations] = useState>({}); const [coverage, setCoverage] = useState(null); const [acceptance, setAcceptance] = useState(null); const [deepAcceptance, setDeepAcceptance] = useState(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(`/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(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(`/api/jobs/${job.id}`); setSelectedJob(detail); @@ -565,13 +613,18 @@ function App() {

Files

数据集浏览

- +
+ + +
{datasets.map((dataset) => ( +
{ @@ -589,6 +642,10 @@ function App() { {dataset.name} {dataset.counts.images} image · {dataset.counts.labels} label · {dataset.counts.masks} mask
+
+ + YOLO {datasetValidations[dataset.name]?.pairs.image_label ?? 0} pair · Mask {datasetValidations[dataset.name]?.pairs.image_mask ?? 0} pair +
{["images", "labels", "masks"].flatMap((kind) => (dataset.samples[kind] ?? []).slice(0, 4).map((sample) => ( @@ -600,8 +657,10 @@ function App() { )}
+
))} + {selectedValidation && } @@ -819,6 +878,31 @@ function ResultPreview({ results }: { results: ResultItem[] }) { ); } +function DatasetQuality({ validation }: { validation: DatasetValidation }) { + return ( +
+
+ {validation.dataset} + {validation.ready.yolo ? "YOLO READY" : validation.ready.mask ? "MASK READY" : "CHECK"} +
+
+
Image/Label{validation.pairs.image_label}
+
Image/Mask{validation.pairs.image_mask}
+
Classes{validation.classes.length || 0}
+
Annotations{validation.counts.annotations}
+
+
+ {validation.checks.map((check) => ( +
+ {check.name} + {check.passed ? "ok" : `${check.errors?.length ?? 0} issue`} +
+ ))} +
+
+ ); +} + function CurvePanel({ curves, selected, diff --git a/frontend/src/styles.css b/frontend/src/styles.css index d2fcf02..ae9cdb4 100644 --- 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)); } }