296 lines
12 KiB
Python
296 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from .config import settings
|
|
from .paths import rel
|
|
|
|
|
|
SEGMODEL_ARCHS = [
|
|
"Unet",
|
|
"UnetPlusPlus",
|
|
"FPN",
|
|
"PSPNet",
|
|
"DeepLabV3",
|
|
"DeepLabV3Plus",
|
|
"Linknet",
|
|
"MAnet",
|
|
"PAN",
|
|
"UPerNet",
|
|
"Segformer",
|
|
"DPT",
|
|
]
|
|
|
|
|
|
YOLO_MODELS = [
|
|
"YOLOv8n-seg",
|
|
"YOLOv8s-seg",
|
|
"YOLOv8m-seg",
|
|
"YOLOv8l-seg",
|
|
"YOLOv8x-seg",
|
|
"YOLOv9c-seg",
|
|
"YOLOv9e-seg",
|
|
"YOLO11n-seg",
|
|
"YOLO11s-seg",
|
|
"YOLO11m-seg",
|
|
"YOLO11l-seg",
|
|
"YOLO11x-seg",
|
|
"YOLO12-seg",
|
|
]
|
|
|
|
|
|
TASK_TYPES = [
|
|
"mock.echo",
|
|
"system.backup",
|
|
"system.check_graph_card",
|
|
"dataset.rename",
|
|
"dataset.to_png",
|
|
"dataset.resize",
|
|
"dataset.resize_single",
|
|
"dataset.pair",
|
|
"dataset.rebuild_labels",
|
|
"dataset.deal_labels",
|
|
"dataset.deal_labels_old",
|
|
"dataset.stack",
|
|
"dataset.stack_single",
|
|
"dataset.stitch",
|
|
"dataset.stitch_single",
|
|
"dataset.run_wizard",
|
|
"dataset.stack_pair_check",
|
|
"dataset.stack_tool_batch",
|
|
"dataset.stack_tool_single",
|
|
"dataset.video_frames",
|
|
"dataset.yolo_check_pairs",
|
|
"dataset.yolo_stack",
|
|
"dataset.yolo_stack_single",
|
|
"dataset.yolo_rebuild_labels",
|
|
"dataset.yolo_txt_ori",
|
|
"dataset.yolo_txt_sort",
|
|
"dataset.yolo_convert_png",
|
|
"dataset.yolo_resize",
|
|
"segmodel.train",
|
|
"segmodel.batch_train",
|
|
"segmodel.predict",
|
|
"segmodel.batch_predict",
|
|
"segmodel.flops",
|
|
"segmodel.params_flops",
|
|
"segmodel.benchmark",
|
|
"segmodel.raw_mask_check",
|
|
"segmodel.metrics",
|
|
"segmodel.copy_best",
|
|
"yolo.train",
|
|
"yolo.train_custom",
|
|
"yolo.batch_train",
|
|
"yolo.predict",
|
|
"yolo.predict_custom",
|
|
"yolo.predict_v1",
|
|
"yolo.batch_predict",
|
|
"yolo.heatmap",
|
|
"yolo.heatmap_custom",
|
|
"yolo.compare",
|
|
"yolo.raw_mask_check",
|
|
"yolo.copy_best",
|
|
"yolo.video_visible",
|
|
"yolo.video_unvisible",
|
|
"yolo.layer_tester",
|
|
"mmseg.init_weights",
|
|
"mmseg.generate_data",
|
|
"mmseg.generate_data_v1",
|
|
"mmseg.generate_data_legacy",
|
|
"mmseg.generate_alg",
|
|
"mmseg.generate_alg_v1",
|
|
"mmseg.generate_alg_legacy",
|
|
"mmseg.train",
|
|
"mmseg.metrics",
|
|
"mmseg.metrics_v1",
|
|
"mmseg.flops_fps",
|
|
"mmseg.flops_fps_v1",
|
|
"mmseg.draw",
|
|
"mmseg.extract_loss_miou",
|
|
"mmseg.delete_epoch",
|
|
"mmseg.copy_result",
|
|
"mmseg.predict_v1",
|
|
"mmseg.predict_v2",
|
|
"visual.train",
|
|
"visual.inference",
|
|
"visual.fps",
|
|
"visual.yolo11_heatmap_v1",
|
|
"visual.yolo11_heatmap_v2",
|
|
"visual.label_ori",
|
|
"visual.label_sort",
|
|
"visual.gen_8bit_png",
|
|
"visual.deal_labels",
|
|
"visual.tool_deal_labels_demo",
|
|
"analysis.all",
|
|
]
|
|
|
|
COMMON_LABEL_DEFAULTS = {
|
|
"src_label_fold": "./Label",
|
|
"save_pro_label_fold": "./ORI_pro_label_fold",
|
|
"save_GT_label_fold": "./ORI_GT_label_fold",
|
|
"Label_Max_Search_layer": 1000,
|
|
"pro_append_name": "_label",
|
|
"GT_append_name": "",
|
|
"GT_channel": 1,
|
|
"back_gnd_color": 0,
|
|
"first_class_color": 1,
|
|
"pic_type": "png",
|
|
"Max_width": 10000,
|
|
"Rebuild_from": "label",
|
|
"Rebuild_to": "GT",
|
|
"save_process_pics": False,
|
|
}
|
|
|
|
|
|
TASK_DEFAULTS: dict[str, dict[str, Any]] = {
|
|
"mock.echo": {"message": "hello from Seg Data Server"},
|
|
"system.backup": {},
|
|
"system.check_graph_card": {},
|
|
"dataset.rename": {"image_dir": "../DataSet_Own/ORI", "label_dir": "../DataSet_Own/Label"},
|
|
"dataset.to_png": {"input_dir": "../DataSet_Own/ORI", "output_dir": "../DataSet_Own/ORI_PNG"},
|
|
"dataset.resize": {"image_dir": "../DataSet_Own/ORI", "label_dir": "../DataSet_Own/Label", "width": 1920, "height": 1080},
|
|
"dataset.resize_single": {"folder": "../DataSet_Own/ORI", "nearest": False, "width": 1920, "height": 1080},
|
|
"dataset.pair": {"image_dir": "../DataSet_Own/ORI", "label_dir": "../DataSet_Own/Label", "prefix": "", "suffix": ""},
|
|
"dataset.rebuild_labels": {"label_dir": "../DataSet_Own/Label"},
|
|
"dataset.deal_labels": COMMON_LABEL_DEFAULTS,
|
|
"dataset.deal_labels_old": COMMON_LABEL_DEFAULTS,
|
|
"dataset.stack": {"image_dir": "../DataSet_Own/ORI", "label_dir": "../DataSet_Own/Label", "result_dir": "../DataSet_Own/stacked", "alpha": 0.3},
|
|
"dataset.stack_single": {"image_path": "../DataSet_Own/ORI/example.png", "label_path": "../DataSet_Own/Label/example.png", "result_dir": "../DataSet_Own/stacked", "alpha": 0.3},
|
|
"dataset.stitch": {"image_dir": "../DataSet_Own/ORI", "label_dir": "../DataSet_Own/Label", "result_dir": "../DataSet_Own/stitch"},
|
|
"dataset.stitch_single": {"image_path": "../DataSet_Own/ORI/example.png", "label_path": "../DataSet_Own/Label/example.png", "result_dir": "../DataSet_Own/stitch", "relative_pos": "up_down"},
|
|
"dataset.run_wizard": {"image_dir": "../DataSet_Own/ORI", "label_dir": "../DataSet_Own/Label", "stdin_text": "7\n"},
|
|
"dataset.stack_pair_check": {"image_dir": "../DataSet_Own/ORI", "label_dir": "../DataSet_Own/Label"},
|
|
"dataset.stack_tool_batch": {"image_dir": "../DataSet_Own/ORI", "label_dir": "../DataSet_Own/Label", "result_dir": "../DataSet_Own/stacked", "alpha": 0.3},
|
|
"dataset.stack_tool_single": {"image_path": "../DataSet_Own/ORI/example.png", "label_path": "../DataSet_Own/Label/example.png", "result_dir": "../DataSet_Own/stacked", "alpha": 0.3},
|
|
"dataset.video_frames": {"video": "../Seg_Predict_Own_Video_V2/LC_Video_1.mp4", "interval": 0.5, "resize": "1920x1080"},
|
|
"dataset.yolo_check_pairs": {"image_dir": "../Seg_All_In_One_YoloModel/Yolo数据集构建/ORI", "label_dir": "../Seg_All_In_One_YoloModel/Yolo数据集构建/Label", "yes": False},
|
|
"dataset.yolo_stack": {"image_dir": "../Seg_All_In_One_YoloModel/Yolo数据集构建/ORI", "label_dir": "../Seg_All_In_One_YoloModel/Yolo数据集构建/Label", "result_dir": "../Seg_All_In_One_YoloModel/Yolo数据集构建/stacked", "alpha": 0.3},
|
|
"dataset.yolo_stack_single": {"image_path": "../Seg_All_In_One_YoloModel/Yolo数据集构建/ORI/example.png", "label_path": "../Seg_All_In_One_YoloModel/Yolo数据集构建/Label/example.png", "result_dir": "../Seg_All_In_One_YoloModel/Yolo数据集构建/stacked", "alpha": 0.3},
|
|
"dataset.yolo_rebuild_labels": COMMON_LABEL_DEFAULTS,
|
|
"dataset.yolo_txt_ori": {},
|
|
"dataset.yolo_txt_sort": {},
|
|
"dataset.yolo_convert_png": {"folder": "../Seg_All_In_One_YoloModel/Yolo数据集构建/Data/images/train", "delete_source": False, "workers": 4},
|
|
"dataset.yolo_resize": {"folder": "../Seg_All_In_One_YoloModel/Yolo数据集构建/Data/images/train", "size": 1080, "workers": 4},
|
|
"segmodel.train": {"architecture": "Unet"},
|
|
"segmodel.batch_train": {},
|
|
"segmodel.predict": {"architecture": "Unet", "run_choice": 1},
|
|
"segmodel.batch_predict": {},
|
|
"segmodel.flops": {"script": "2_predict_params_and_FLOPs_V2.py"},
|
|
"segmodel.params_flops": {"architecture": "Unet", "shape": [512, 512]},
|
|
"segmodel.benchmark": {"architecture": "Unet", "shape": [512, 512], "repeat_times": 3},
|
|
"segmodel.raw_mask_check": {},
|
|
"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_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.predict_v1": {"model": "YOLOv8n-seg", "pt_name": "best.pt", "conf": 0.2, "run_choice": 1},
|
|
"yolo.batch_predict": {"pt_name": "best.pt", "conf": 0.2},
|
|
"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"},
|
|
"yolo.compare": {"pt_name": "all"},
|
|
"yolo.raw_mask_check": {"pt_name": "best.pt"},
|
|
"yolo.copy_best": {"pt_name": "best.pt"},
|
|
"yolo.video_visible": {},
|
|
"yolo.video_unvisible": {},
|
|
"yolo.layer_tester": {},
|
|
"mmseg.init_weights": {},
|
|
"mmseg.generate_data": {},
|
|
"mmseg.generate_data_v1": {},
|
|
"mmseg.generate_data_legacy": {},
|
|
"mmseg.generate_alg": {"dataset_choice": 1, "gpu_count": 1, "gpu_ids": [0], "schedule_mode": 2, "max_epochs": 300, "algorithm_choice": 1},
|
|
"mmseg.generate_alg_v1": {"dataset_choice": 1, "gpu_count": 1, "gpu_ids": [0], "schedule_mode": 2, "max_epochs": 300, "algorithm_choice": 1},
|
|
"mmseg.generate_alg_legacy": {},
|
|
"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},
|
|
"mmseg.metrics_v1": {"dataset_choice": 1, "algorithm_choice": 0},
|
|
"mmseg.flops_fps": {"input_dir": "../Hardisk", "output_dir": "../BestMode_Predict_Results_DataSet_Public", "repeat_times": 3, "dataset_choice": 1, "algorithm_choice": 0},
|
|
"mmseg.flops_fps_v1": {"dataset_choice": 1, "algorithm_choice": 0},
|
|
"mmseg.draw": {},
|
|
"mmseg.extract_loss_miou": {},
|
|
"mmseg.delete_epoch": {},
|
|
"mmseg.copy_result": {},
|
|
"mmseg.predict_v1": {},
|
|
"mmseg.predict_v2": {},
|
|
"visual.train": {},
|
|
"visual.inference": {},
|
|
"visual.fps": {"weights": "yolov8n.pt", "batch": 1, "imgs": [640, 640], "device": "cpu", "warmup": 10, "testtime": 20, "half": False},
|
|
"visual.yolo11_heatmap_v1": {},
|
|
"visual.yolo11_heatmap_v2": {},
|
|
"visual.label_ori": {},
|
|
"visual.label_sort": {},
|
|
"visual.gen_8bit_png": {},
|
|
"visual.deal_labels": COMMON_LABEL_DEFAULTS,
|
|
"visual.tool_deal_labels_demo": {},
|
|
"analysis.all": {"input_dir": "../BestMode_Predict_Results_DataSet_Public", "output_dir": "./", "dataset_choice": 1},
|
|
}
|
|
|
|
|
|
def _read_json(path: Path) -> Any | None:
|
|
try:
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def discover_datasets() -> list[dict[str, Any]]:
|
|
root = settings.source_root
|
|
candidates: list[dict[str, Any]] = []
|
|
for base in ["DataSet_Public", "BestMode_Predict_Results_DataSet_Public", "Hardisk"]:
|
|
parent = root / base
|
|
if not parent.exists():
|
|
continue
|
|
for item in sorted(parent.iterdir()):
|
|
if item.is_dir():
|
|
candidates.append({"name": item.name, "path": rel(item, root), "source": base})
|
|
mmseg_params = root / "Seg_All_In_One_MMSeg" / "My_All_In_One" / "1_Data_Parameter"
|
|
for item in sorted(mmseg_params.glob("*.json")):
|
|
data = _read_json(item)
|
|
if item.name == "All_Data_Record.json" or not data:
|
|
continue
|
|
candidates.append({"name": item.stem, "path": rel(item, root), "source": "mmseg_parameter"})
|
|
uploaded_root = settings.project_root / "var" / "uploads" / "datasets"
|
|
if uploaded_root.exists():
|
|
for item in sorted(uploaded_root.iterdir()):
|
|
if item.is_dir():
|
|
candidates.append({"name": item.name, "path": rel(item, settings.project_root), "source": "uploaded"})
|
|
return candidates
|
|
|
|
|
|
def discover_mmseg_algorithms() -> list[str]:
|
|
alg_dir = settings.source_root / "Seg_All_In_One_MMSeg" / "My_All_In_One" / "2_Alg_Program"
|
|
if not alg_dir.exists():
|
|
return []
|
|
return sorted(path.stem for path in alg_dir.glob("*.py"))
|
|
|
|
|
|
def discover_weights_summary() -> dict[str, Any]:
|
|
manifest = settings.weights_root / "manifest.json"
|
|
if not manifest.exists():
|
|
return {"manifest": None, "count": 0, "total_bytes": 0}
|
|
data = _read_json(manifest) or {}
|
|
return {
|
|
"manifest": rel(manifest, settings.project_root),
|
|
"count": len(data.get("files", [])),
|
|
"total_bytes": data.get("total_bytes", 0),
|
|
"updated_at": data.get("updated_at"),
|
|
}
|
|
|
|
|
|
def get_catalog() -> dict[str, Any]:
|
|
return {
|
|
"source_root": str(settings.source_root),
|
|
"project_root": str(settings.project_root),
|
|
"task_types": TASK_TYPES,
|
|
"task_defaults": TASK_DEFAULTS,
|
|
"segmodel_architectures": SEGMODEL_ARCHS,
|
|
"yolo_models": YOLO_MODELS,
|
|
"mmseg_algorithms": discover_mmseg_algorithms(),
|
|
"datasets": discover_datasets(),
|
|
"weights": discover_weights_summary(),
|
|
}
|