131 lines
6.4 KiB
Python
131 lines
6.4 KiB
Python
from __future__ import annotations
|
|
|
|
from ...commands import CommandSpec, append_flag, bash, conda_python, required
|
|
from ...config import settings
|
|
|
|
|
|
MMSEG_DIR = settings.source_root / "Seg_All_In_One_MMSeg"
|
|
MY_DIR = MMSEG_DIR / "My_All_In_One"
|
|
|
|
|
|
def _stdin_for_generate_alg(params: dict) -> str:
|
|
lines = [
|
|
str(params.get("dataset_choice", 1)),
|
|
str(params.get("gpu_count", 1)),
|
|
]
|
|
gpu_ids = params.get("gpu_ids", [0])
|
|
if isinstance(gpu_ids, str):
|
|
gpu_ids = [part.strip() for part in gpu_ids.split(",") if part.strip()]
|
|
for index in range(int(params.get("gpu_count", len(gpu_ids) or 1))):
|
|
lines.append(str(gpu_ids[index] if index < len(gpu_ids) else 0))
|
|
|
|
mode = str(params.get("schedule_mode", 2))
|
|
lines.append(mode)
|
|
if mode == "1":
|
|
lines.extend(
|
|
[
|
|
str(params.get("train_k", 40)),
|
|
str(params.get("check_count", 10)),
|
|
str(params.get("logger_interval", 50)),
|
|
]
|
|
)
|
|
else:
|
|
lines.extend(
|
|
[
|
|
str(params.get("max_epochs", 300)),
|
|
str(params.get("val_interval", 1)),
|
|
str(params.get("checkpoint_interval", 10)),
|
|
str(params.get("logger_interval", "")),
|
|
]
|
|
)
|
|
lines.append(str(params.get("algorithm_choice", 1)))
|
|
return "\n".join(lines) + "\n"
|
|
|
|
|
|
def build_mmseg_task(job_type: str, params: dict, conda_env: str) -> CommandSpec | None:
|
|
if job_type == "mmseg.init_weights":
|
|
return CommandSpec(conda_python(conda_env, MY_DIR / "0_Initial_Save_All_Model_locally.py"), MMSEG_DIR, "download/save MMSeg pretrained weights locally")
|
|
|
|
if job_type == "mmseg.generate_data":
|
|
return CommandSpec(conda_python(conda_env, MY_DIR / "1_Initial_Data_All_data_from_1_Data_Parameter-V2.py"), MMSEG_DIR, "generate MMSeg dataset configs from JSON parameters")
|
|
|
|
if job_type == "mmseg.generate_data_v1":
|
|
return CommandSpec(conda_python(conda_env, MY_DIR / "1_Initial_Data_All_data_from_1_Data_Parameter-V1.py"), MMSEG_DIR, "generate MMSeg dataset configs with V1 flow")
|
|
|
|
if job_type == "mmseg.generate_data_legacy":
|
|
return CommandSpec(conda_python(conda_env, MY_DIR / "1_Initial_Data_All-ori.py"), MMSEG_DIR, "run legacy original MMSeg dataset config generator")
|
|
|
|
if job_type == "mmseg.generate_alg":
|
|
script = MY_DIR / "2_Initial_Alg_All_data_from_2_Alg_Program-V2.py"
|
|
return CommandSpec(
|
|
conda_python(conda_env, script),
|
|
MMSEG_DIR,
|
|
"generate MMSeg algorithm config and training command",
|
|
stdin_text=_stdin_for_generate_alg(params),
|
|
)
|
|
|
|
if job_type == "mmseg.generate_alg_v1":
|
|
script = MY_DIR / "2_Initial_Alg_All_data_from_2_Alg_Program-V1.py"
|
|
return CommandSpec(
|
|
conda_python(conda_env, script),
|
|
MMSEG_DIR,
|
|
"generate MMSeg algorithm config and training command with V1 flow",
|
|
stdin_text=_stdin_for_generate_alg(params),
|
|
)
|
|
|
|
if job_type == "mmseg.generate_alg_legacy":
|
|
return CommandSpec(conda_python(conda_env, MY_DIR / "2_Initial_Alg_All-ori-old.py"), MMSEG_DIR, "run legacy original MMSeg algorithm config generator")
|
|
|
|
if job_type == "mmseg.train":
|
|
config_path = required(params, "config")
|
|
args = conda_python(conda_env, MMSEG_DIR / "tools" / "train.py", config_path)
|
|
append_flag(args, "--work-dir", params.get("work_dir"))
|
|
return CommandSpec(args, MMSEG_DIR, "train MMSeg model")
|
|
|
|
if job_type == "mmseg.metrics":
|
|
args = conda_python(conda_env, MY_DIR / "4_2_predict_matrics_from_log_V2.py")
|
|
append_flag(args, "--input_dir", params.get("input_dir", "../Hardisk"))
|
|
append_flag(args, "--output_dir", params.get("output_dir", "../BestMode_Predict_Results_DataSet_Public"))
|
|
stdin = f"{params.get('dataset_choice', 1)}\n{params.get('algorithm_choice', 0)}\n"
|
|
return CommandSpec(args, MMSEG_DIR, "extract best MMSeg metrics from logs", stdin_text=stdin)
|
|
|
|
if job_type == "mmseg.metrics_v1":
|
|
args = conda_python(conda_env, MY_DIR / "4_2_predict_matrics_from_log_V1.py")
|
|
stdin = f"{params.get('dataset_choice', 1)}\n{params.get('algorithm_choice', 0)}\n"
|
|
return CommandSpec(args, MMSEG_DIR, "extract best MMSeg metrics from logs with V1 script", stdin_text=stdin)
|
|
|
|
if job_type == "mmseg.flops_fps":
|
|
args = conda_python(conda_env, MY_DIR / "4_1_predict_params_FLOPs_FPS_V2.py")
|
|
append_flag(args, "--input_dir", params.get("input_dir", "../Hardisk"))
|
|
append_flag(args, "--output_dir", params.get("output_dir", "../BestMode_Predict_Results_DataSet_Public"))
|
|
append_flag(args, "--repeat-times", params.get("repeat_times", 3))
|
|
stdin = f"{params.get('dataset_choice', 1)}\n{params.get('algorithm_choice', 0)}\n"
|
|
if "shape_h" in params and "shape_w" in params:
|
|
stdin += f"{params['shape_h']}\n{params['shape_w']}\n"
|
|
return CommandSpec(args, MMSEG_DIR, "calculate MMSeg FLOPs/Params/FPS", stdin_text=stdin)
|
|
|
|
if job_type == "mmseg.flops_fps_v1":
|
|
args = conda_python(conda_env, MY_DIR / "4_1_predict_params_and_FLOPs_V1.py")
|
|
stdin = f"{params.get('dataset_choice', 1)}\n{params.get('algorithm_choice', 0)}\n"
|
|
return CommandSpec(args, MMSEG_DIR, "calculate MMSeg FLOPs/Params/FPS with V1 script", stdin_text=stdin)
|
|
|
|
if job_type == "mmseg.draw":
|
|
return CommandSpec(conda_python(conda_env, MY_DIR / "4_3_predict_draw_pictures_and_tabels.py"), MMSEG_DIR, "generate MMSeg prediction pictures and tables")
|
|
|
|
if job_type == "mmseg.extract_loss_miou":
|
|
return CommandSpec(conda_python(conda_env, MY_DIR / "4_4_extract_loss_and_best_miou.py"), MMSEG_DIR, "extract MMSeg loss and best mIoU curves")
|
|
|
|
if job_type == "mmseg.delete_epoch":
|
|
return CommandSpec(conda_python(conda_env, MY_DIR / "3_Find_And_Delete_Special_Epoch.py"), MMSEG_DIR, "find and delete selected epoch checkpoints")
|
|
|
|
if job_type == "mmseg.copy_result":
|
|
return CommandSpec(bash(MY_DIR / "3_Tool_Copy_Result_To_Hardisk.sh"), MMSEG_DIR, "copy MMSeg results to Hardisk")
|
|
|
|
if job_type == "mmseg.predict_v1":
|
|
return CommandSpec(conda_python(conda_env, MY_DIR / "x4_Predict_V1-.py"), MMSEG_DIR, "run MMSeg prediction V1 helper")
|
|
|
|
if job_type == "mmseg.predict_v2":
|
|
return CommandSpec(conda_python(conda_env, MY_DIR / "x4_Predict_V2-.py"), MMSEG_DIR, "run MMSeg prediction V2 helper")
|
|
|
|
return None
|