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_v1", "yolo.batch_predict", "yolo.heatmap", "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_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.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(), }