141 lines
3.5 KiB
Python
141 lines
3.5 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",
|
|
"dataset.rename",
|
|
"dataset.to_png",
|
|
"dataset.resize",
|
|
"dataset.pair",
|
|
"dataset.rebuild_labels",
|
|
"dataset.stack",
|
|
"dataset.stitch",
|
|
"dataset.video_frames",
|
|
"segmodel.train",
|
|
"segmodel.batch_train",
|
|
"segmodel.predict",
|
|
"segmodel.batch_predict",
|
|
"segmodel.flops",
|
|
"segmodel.raw_mask_check",
|
|
"segmodel.metrics",
|
|
"yolo.train",
|
|
"yolo.batch_train",
|
|
"yolo.predict",
|
|
"yolo.batch_predict",
|
|
"yolo.heatmap",
|
|
"yolo.compare",
|
|
"yolo.raw_mask_check",
|
|
"yolo.video_visible",
|
|
"yolo.video_unvisible",
|
|
"mmseg.init_weights",
|
|
"mmseg.generate_data",
|
|
"mmseg.generate_alg",
|
|
"mmseg.train",
|
|
"mmseg.metrics",
|
|
"mmseg.flops_fps",
|
|
"mmseg.draw",
|
|
"mmseg.extract_loss_miou",
|
|
"analysis.all",
|
|
]
|
|
|
|
|
|
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"})
|
|
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,
|
|
"segmodel_architectures": SEGMODEL_ARCHS,
|
|
"yolo_models": YOLO_MODELS,
|
|
"mmseg_algorithms": discover_mmseg_algorithms(),
|
|
"datasets": discover_datasets(),
|
|
"weights": discover_weights_summary(),
|
|
}
|
|
|