70 lines
3.4 KiB
Python
70 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
from ...commands import CommandSpec, append_flag, bash, conda_python, required
|
|
from ...config import settings
|
|
|
|
|
|
SEGMODEL_DIR = settings.source_root / "Seg_All_In_One_SegModel"
|
|
|
|
|
|
def build_segmodel_task(job_type: str, params: dict, conda_env: str) -> CommandSpec | None:
|
|
env = {"SEG_CONDA_ENV": conda_env}
|
|
|
|
if job_type == "segmodel.train":
|
|
args = conda_python(conda_env, SEGMODEL_DIR / "train.py")
|
|
append_flag(args, "-a", required(params, "architecture"))
|
|
return CommandSpec(args, SEGMODEL_DIR, "train one segmentation_models_pytorch architecture")
|
|
|
|
if job_type == "segmodel.batch_train":
|
|
return CommandSpec(bash(SEGMODEL_DIR / "train.sh"), SEGMODEL_DIR, "run legacy SegModel batch training", env=env)
|
|
|
|
if job_type == "segmodel.predict":
|
|
args = conda_python(conda_env, SEGMODEL_DIR / "1_predict.py")
|
|
append_flag(args, "-a", required(params, "architecture"))
|
|
choice = str(params.get("run_choice", 1))
|
|
return CommandSpec(args, SEGMODEL_DIR, "predict with one SegModel run", stdin_text=f"{choice}\n")
|
|
|
|
if job_type == "segmodel.batch_predict":
|
|
return CommandSpec(bash(SEGMODEL_DIR / "predict.sh"), SEGMODEL_DIR, "run legacy SegModel batch prediction", env=env)
|
|
|
|
if job_type == "segmodel.flops":
|
|
script = SEGMODEL_DIR / params.get("script", "2_predict_params_and_FLOPs_V2.py")
|
|
return CommandSpec(conda_python(conda_env, script), SEGMODEL_DIR, "calculate SegModel params/FLOPs/FPS")
|
|
|
|
if job_type == "segmodel.params_flops":
|
|
args = conda_python(conda_env, SEGMODEL_DIR / "Tool_get_params_and_FLOPs.py")
|
|
append_flag(args, "-a", required(params, "architecture"))
|
|
shape = params.get("shape", [512, 512])
|
|
if isinstance(shape, str):
|
|
shape = [part for part in shape.replace(",", " ").split() if part]
|
|
if shape:
|
|
args.append("--shape")
|
|
args.extend(str(item) for item in shape)
|
|
return CommandSpec(args, SEGMODEL_DIR, "calculate SegModel params and FLOPs for one architecture")
|
|
|
|
if job_type == "segmodel.benchmark":
|
|
args = conda_python(conda_env, SEGMODEL_DIR / "Tool_benchmark_smp.py")
|
|
append_flag(args, "-a", required(params, "architecture"))
|
|
append_flag(args, "-c", params.get("checkpoint"))
|
|
shape = params.get("shape", [512, 512])
|
|
if isinstance(shape, str):
|
|
shape = [part for part in shape.replace(",", " ").split() if part]
|
|
if shape:
|
|
args.append("--shape")
|
|
args.extend(str(item) for item in shape)
|
|
append_flag(args, "--repeat-times", params.get("repeat_times"))
|
|
append_flag(args, "--log-interval", params.get("log_interval"))
|
|
stdin = None if params.get("checkpoint") else params.get("stdin_text")
|
|
return CommandSpec(args, SEGMODEL_DIR, "benchmark SegModel FPS/latency", stdin_text=stdin)
|
|
|
|
if job_type == "segmodel.raw_mask_check":
|
|
return CommandSpec(conda_python(conda_env, SEGMODEL_DIR / "1_predict_raw_masks_check.py"), SEGMODEL_DIR, "check SegModel raw mask completeness")
|
|
|
|
if job_type == "segmodel.metrics":
|
|
return CommandSpec(conda_python(conda_env, SEGMODEL_DIR / "3_predict_matrics_from_log.py"), SEGMODEL_DIR, "parse SegModel training/prediction metrics")
|
|
|
|
if job_type == "segmodel.copy_best":
|
|
return CommandSpec(bash(SEGMODEL_DIR / "Tool_Copy_Best_Model.sh"), SEGMODEL_DIR, "copy best SegModel weights to prediction area")
|
|
|
|
return None
|