admin/Seg_Data_Server_Net
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add custom YOLO prediction and heatmap workflow

Browse Source
This commit is contained in:
admin
2026-06-30 15:11:47 +08:00
parent 4d0c26be05
commit 777f168a75
12 changed files with 393 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

90
frontend/src/main.tsx
View File

@@ -276,7 +276,9 @@ const defaultParams: Record<string, Record<string, unknown>> = {
"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.predict_custom": { weights: "var/custom_yolo_runs/example/weights/best.pt", source: "var/uploads/datasets/example/images", imgsz: 640, conf: 0.25, device: "cpu", name: "example_predict", exist_ok: true },
"yolo.heatmap": { model: "YOLOv8n-seg", cam_method: "All", pt_name: "best.pt", run_choice: 1 },
"yolo.heatmap_custom": { weights: "var/custom_yolo_runs/example/weights/best.pt", source: "var/uploads/datasets/example/images", model_key: "YOLO11n-seg", cam_method: "GradCAM", target_layers: "model.model.model[9]", limit: 3, name: "example_heatmap" },
"mmseg.generate_alg": { dataset_choice: 1, gpu_count: 1, gpu_ids: [0], schedule_mode: 2, max_epochs: 300, algorithm_choice: 1 },
"mmseg.train": { config: "configs/example.py", work_dir: "../DataSet_Public_outputs/example" },
"mmseg.metrics": { input_dir: "../Hardisk", output_dir: "../BestMode_Predict_Results_DataSet_Public", dataset_choice: 1, algorithm_choice: 0 },
@@ -558,17 +560,84 @@ function App() {
}
}
function customYoloWeightPath(dataset: UploadedDataset) {
const expected = `var/custom_yolo_runs/${dataset.name}/weights/best.pt`;
return results.find((item) => item.relative_path === expected || item.relative_path.endsWith(`/custom_yolo_runs/${dataset.name}/weights/best.pt`))?.relative_path ?? expected;
}
async function createSelectedYoloYaml() {
if (!selectedDataset) return;
const classNames = selectedValidation?.classes.map((classId) => `class_${classId}`) ?? undefined;
return api<{ relative_path: string; path: string }>(`/api/datasets/${encodeURIComponent(selectedDataset.name)}/yolo-yaml`, {
method: "POST",
body: JSON.stringify({ class_names: classNames })
});
}
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 })
});
const generated = await createSelectedYoloYaml();
if (!generated) return;
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));
setParams(JSON.stringify({ model: "YOLO11n-seg", data: generated.path, epochs: 10, imgsz: 640, batch: 1, workers: 0, device: "cpu", project: "var/custom_yolo_runs", name: selectedDataset.name, exist_ok: true }, null, 2));
await refresh();
} finally {
setBusy(false);
}
}
async function startSelectedYoloTrain() {
if (!selectedDataset) return;
setBusy(true);
try {
const generated = await createSelectedYoloYaml();
if (!generated) return;
await api<Job>("/api/jobs", {
method: "POST",
body: JSON.stringify({
type: "yolo.train_custom",
params: { model: "YOLO11n-seg", data: generated.path, epochs: 10, imgsz: 640, batch: 1, workers: 0, device: "cpu", project: "var/custom_yolo_runs", name: selectedDataset.name, exist_ok: true }
})
});
window.location.hash = "jobs";
await refresh();
} finally {
setBusy(false);
}
}
async function startSelectedYoloPredict() {
if (!selectedDataset?.absolute_layout) return;
setBusy(true);
try {
await api<Job>("/api/jobs", {
method: "POST",
body: JSON.stringify({
type: "yolo.predict_custom",
params: { weights: customYoloWeightPath(selectedDataset), source: selectedDataset.absolute_layout.images, imgsz: 640, conf: 0.25, device: "cpu", project: "var/custom_yolo_runs", name: `${selectedDataset.name}_predict`, exist_ok: true }
})
});
window.location.hash = "jobs";
await refresh();
} finally {
setBusy(false);
}
}
async function startSelectedYoloHeatmap() {
if (!selectedDataset?.absolute_layout) return;
setBusy(true);
try {
await api<Job>("/api/jobs", {
method: "POST",
body: JSON.stringify({
type: "yolo.heatmap_custom",
params: { weights: customYoloWeightPath(selectedDataset), source: selectedDataset.absolute_layout.images, model_key: "YOLO11n-seg", cam_method: "GradCAM", target_layers: "model.model.model[9]", limit: 3, project: "var/custom_yolo_runs", name: `${selectedDataset.name}_heatmap` }
})
});
window.location.hash = "jobs";
await refresh();
} finally {
setBusy(false);
@@ -807,6 +876,15 @@ function App() {
<button className="iconButton" disabled={busy || !selectedValidation?.ready.yolo} onClick={generateSelectedYoloYaml} title="生成 YOLO dataset.yaml">
<FileSearch size={18} />
</button>
<button className="iconButton" disabled={busy || !selectedValidation?.ready.yolo} onClick={startSelectedYoloTrain} title="启动自定义 YOLO 训练">
<Play size={18} />
</button>
<button className="iconButton" disabled={busy || !selectedDataset?.absolute_layout} onClick={startSelectedYoloPredict} title="使用自定义 best.pt 预测">
<FileImage size={18} />
</button>
<button className="iconButton" disabled={busy || !selectedDataset?.absolute_layout} onClick={startSelectedYoloHeatmap} title="使用自定义 best.pt 生成热度图">
<Zap size={18} />
</button>
<FileImage size={22} />
</div>
</div>