230 lines
11 KiB
Python
230 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
from ...commands import CommandSpec, append_flag, bash, conda_python, required
|
|
from ...config import settings
|
|
|
|
|
|
DATASET_TOOL_DIR = settings.source_root / "DataSet_Own" / "1. 图片预处理（内含使用手册）"
|
|
STACK_TOOL_DIR = settings.source_root / "Tool-图片堆叠"
|
|
VIDEO_DIR = settings.source_root / "Seg_Predict_Own_Video_V2"
|
|
YOLO_DATASET_DIR = settings.source_root / "Seg_All_In_One_YoloModel" / "Yolo数据集构建"
|
|
|
|
|
|
def _dataset_script(name: str) -> Path:
|
|
return DATASET_TOOL_DIR / name
|
|
|
|
|
|
def _append_label_rebuild_flags(args: list[str], params: dict) -> None:
|
|
append_flag(args, "-src_fold", params.get("src_label_fold"))
|
|
append_flag(args, "-save_pro_fold", params.get("save_pro_label_fold"))
|
|
append_flag(args, "-save_GT_fold", params.get("save_GT_label_fold"))
|
|
append_flag(args, "-fold_search_depth", params.get("Label_Max_Search_layer"))
|
|
append_flag(args, "-pro_suffix_name", params.get("pro_append_name"))
|
|
append_flag(args, "-GT_suffix_name", params.get("GT_append_name"))
|
|
append_flag(args, "-GT_channel", params.get("GT_channel"))
|
|
append_flag(args, "-back_gnd_color", params.get("back_gnd_color"))
|
|
append_flag(args, "-first_class_color", params.get("first_class_color"))
|
|
append_flag(args, "-pic_type", params.get("pic_type"))
|
|
append_flag(args, "-Max_width", params.get("Max_width"))
|
|
append_flag(args, "-Rebuild_from", params.get("Rebuild_from"))
|
|
append_flag(args, "-Rebuild_to", params.get("Rebuild_to"))
|
|
if "save_process_pics" in params:
|
|
append_flag(args, "-save_process_pics", str(params["save_process_pics"]))
|
|
|
|
|
|
def build_dataset_task(job_type: str, params: dict, conda_env: str) -> CommandSpec | None:
|
|
if job_type == "dataset.rename":
|
|
args = bash(_dataset_script("1_rename_pics.sh"))
|
|
append_flag(args, "-i", required(params, "image_dir"))
|
|
append_flag(args, "-l", required(params, "label_dir"))
|
|
return CommandSpec(args, DATASET_TOOL_DIR, "rename and normalize image/label names")
|
|
|
|
if job_type == "dataset.to_png":
|
|
script = _dataset_script("2_1_Trans_to_png.py")
|
|
args = conda_python(conda_env, script)
|
|
append_flag(args, "-i", params.get("input_dir"))
|
|
append_flag(args, "-o", params.get("output_dir"))
|
|
return CommandSpec(args, DATASET_TOOL_DIR, "convert images to png")
|
|
|
|
if job_type == "dataset.resize":
|
|
args = bash(_dataset_script("2_reformate_pics.sh"))
|
|
append_flag(args, "-i", params.get("image_dir"))
|
|
append_flag(args, "-l", params.get("label_dir"))
|
|
append_flag(args, "-w", params.get("width", 1920))
|
|
append_flag(args, "-h", params.get("height", 1080))
|
|
return CommandSpec(args, DATASET_TOOL_DIR, "resize and reformat image/label folders")
|
|
|
|
if job_type == "dataset.resize_single":
|
|
script = _dataset_script("2_2_Resize.py")
|
|
args = conda_python(
|
|
conda_env,
|
|
script,
|
|
required(params, "folder"),
|
|
str(bool(params.get("nearest", False))),
|
|
params.get("width", 1920),
|
|
params.get("height", 1080),
|
|
)
|
|
return CommandSpec(args, DATASET_TOOL_DIR, "resize one image folder with the legacy Python tool")
|
|
|
|
if job_type == "dataset.pair":
|
|
args = bash(_dataset_script("3_pair_ori_label.sh"))
|
|
append_flag(args, "-i", required(params, "image_dir"))
|
|
append_flag(args, "-l", required(params, "label_dir"))
|
|
append_flag(args, "-p", params.get("prefix", ""))
|
|
append_flag(args, "-s", params.get("suffix", ""))
|
|
return CommandSpec(args, DATASET_TOOL_DIR, "check image and label pairing")
|
|
|
|
if job_type == "dataset.rebuild_labels":
|
|
args = bash(_dataset_script("4_rebuild_labels.sh"))
|
|
append_flag(args, "-l", required(params, "label_dir"))
|
|
return CommandSpec(args, DATASET_TOOL_DIR, "rebuild color labels into GT masks")
|
|
|
|
if job_type in {"dataset.deal_labels", "dataset.deal_labels_old"}:
|
|
script_name = "4_deal_labels_old(老版程序).py" if job_type.endswith("_old") else "4_deal_labels.py"
|
|
args = conda_python(conda_env, _dataset_script(script_name))
|
|
_append_label_rebuild_flags(args, params)
|
|
return CommandSpec(args, DATASET_TOOL_DIR, "run legacy label rebuild/conversion script")
|
|
|
|
if job_type == "dataset.stack":
|
|
args = bash(_dataset_script("5_TOOL_stack_pics.sh"))
|
|
append_flag(args, "-i", required(params, "image_dir"))
|
|
append_flag(args, "-l", required(params, "label_dir"))
|
|
append_flag(args, "-r", required(params, "result_dir"))
|
|
append_flag(args, "-a", params.get("alpha", 0.3))
|
|
append_flag(args, "-p", params.get("prefix", ""))
|
|
append_flag(args, "-s", params.get("suffix", ""))
|
|
return CommandSpec(args, DATASET_TOOL_DIR, "overlay image and label for inspection")
|
|
|
|
if job_type == "dataset.stack_single":
|
|
args = conda_python(
|
|
conda_env,
|
|
_dataset_script("5_stack_picture.py"),
|
|
required(params, "image_path"),
|
|
required(params, "label_path"),
|
|
required(params, "result_dir"),
|
|
params.get("alpha", 0.3),
|
|
)
|
|
return CommandSpec(args, DATASET_TOOL_DIR, "overlay one image and label pair")
|
|
|
|
if job_type == "dataset.stitch":
|
|
args = bash(_dataset_script("6_TOOL_stitch_pics.sh"))
|
|
append_flag(args, "-i", required(params, "image_dir"))
|
|
append_flag(args, "-l", required(params, "label_dir"))
|
|
append_flag(args, "-r", required(params, "result_dir"))
|
|
return CommandSpec(args, DATASET_TOOL_DIR, "stitch image and label panels")
|
|
|
|
if job_type == "dataset.stitch_single":
|
|
args = conda_python(
|
|
conda_env,
|
|
_dataset_script("6_stitch_picture.py"),
|
|
required(params, "image_path"),
|
|
required(params, "label_path"),
|
|
required(params, "result_dir"),
|
|
params.get("relative_pos", "up_down"),
|
|
)
|
|
return CommandSpec(args, DATASET_TOOL_DIR, "stitch one image and label pair")
|
|
|
|
if job_type == "dataset.run_wizard":
|
|
args = bash(_dataset_script("Seg_data_run.sh"))
|
|
append_flag(args, "-i", required(params, "image_dir"))
|
|
append_flag(args, "-l", required(params, "label_dir"))
|
|
return CommandSpec(args, DATASET_TOOL_DIR, "run legacy interactive dataset wizard", stdin_text=params.get("stdin_text"))
|
|
|
|
if job_type == "dataset.stack_pair_check":
|
|
args = conda_python(conda_env, STACK_TOOL_DIR / "1_check_picture_pair.py")
|
|
append_flag(args, "-i", required(params, "image_dir"))
|
|
append_flag(args, "-l", required(params, "label_dir"))
|
|
append_flag(args, "-p", params.get("prefix", ""))
|
|
append_flag(args, "-s", params.get("suffix", ""))
|
|
return CommandSpec(args, STACK_TOOL_DIR, "check image/label pair names with stack tool")
|
|
|
|
if job_type == "dataset.stack_tool_batch":
|
|
args = bash(STACK_TOOL_DIR / "2_TOOL_stack_pics.sh")
|
|
append_flag(args, "-i", required(params, "image_dir"))
|
|
append_flag(args, "-l", required(params, "label_dir"))
|
|
append_flag(args, "-r", required(params, "result_dir"))
|
|
append_flag(args, "-a", params.get("alpha", 0.3))
|
|
append_flag(args, "-p", params.get("prefix", ""))
|
|
append_flag(args, "-s", params.get("suffix", ""))
|
|
return CommandSpec(args, STACK_TOOL_DIR, "run standalone stack tool batch overlay")
|
|
|
|
if job_type == "dataset.stack_tool_single":
|
|
args = conda_python(
|
|
conda_env,
|
|
STACK_TOOL_DIR / "2_stack_picture.py",
|
|
required(params, "image_path"),
|
|
required(params, "label_path"),
|
|
required(params, "result_dir"),
|
|
params.get("alpha", 0.3),
|
|
)
|
|
return CommandSpec(args, STACK_TOOL_DIR, "run standalone stack tool for one pair")
|
|
|
|
if job_type == "dataset.video_frames":
|
|
script = VIDEO_DIR / "1_Save_Frame_V2.py"
|
|
args = conda_python(conda_env, script)
|
|
append_flag(args, "--video", required(params, "video"))
|
|
append_flag(args, "--interval", params.get("interval", 0.5))
|
|
append_flag(args, "--resize", params.get("resize"))
|
|
append_flag(args, "--output_dir", params.get("output_dir"))
|
|
return CommandSpec(args, VIDEO_DIR, "extract video frames into DataSet_Public layout")
|
|
|
|
if job_type == "dataset.yolo_check_pairs":
|
|
script = YOLO_DATASET_DIR / "0_1_check_picture_pair.py"
|
|
args = conda_python(conda_env, script)
|
|
append_flag(args, "-i", required(params, "image_dir"))
|
|
append_flag(args, "-l", required(params, "label_dir"))
|
|
append_flag(args, "-p", params.get("prefix", ""))
|
|
append_flag(args, "-s", params.get("suffix", ""))
|
|
if params.get("yes"):
|
|
args.append("-y")
|
|
return CommandSpec(args, YOLO_DATASET_DIR, "check YOLO image/label pairs")
|
|
|
|
if job_type == "dataset.yolo_stack":
|
|
args = bash(YOLO_DATASET_DIR / "0_2_TOOL_stack_pics.sh")
|
|
append_flag(args, "-i", required(params, "image_dir"))
|
|
append_flag(args, "-l", required(params, "label_dir"))
|
|
append_flag(args, "-r", required(params, "result_dir"))
|
|
append_flag(args, "-a", params.get("alpha", 0.3))
|
|
append_flag(args, "-p", params.get("prefix", ""))
|
|
append_flag(args, "-s", params.get("suffix", ""))
|
|
return CommandSpec(args, YOLO_DATASET_DIR, "overlay YOLO dataset image/label pairs")
|
|
|
|
if job_type == "dataset.yolo_stack_single":
|
|
args = conda_python(
|
|
conda_env,
|
|
YOLO_DATASET_DIR / "0_2_stack_picture.py",
|
|
required(params, "image_path"),
|
|
required(params, "label_path"),
|
|
required(params, "result_dir"),
|
|
params.get("alpha", 0.3),
|
|
)
|
|
return CommandSpec(args, YOLO_DATASET_DIR, "overlay one YOLO dataset image/label pair")
|
|
|
|
if job_type == "dataset.yolo_rebuild_labels":
|
|
args = conda_python(conda_env, YOLO_DATASET_DIR / "1_deal_labels.py")
|
|
_append_label_rebuild_flags(args, params)
|
|
return CommandSpec(args, YOLO_DATASET_DIR, "rebuild labels for YOLO dataset generation")
|
|
|
|
if job_type == "dataset.yolo_txt_ori":
|
|
return CommandSpec(conda_python(conda_env, YOLO_DATASET_DIR / "2_Check_and_Gen_Txt_Label_ori_label.py"), YOLO_DATASET_DIR, "generate YOLO txt labels preserving original class ids")
|
|
|
|
if job_type == "dataset.yolo_txt_sort":
|
|
return CommandSpec(conda_python(conda_env, YOLO_DATASET_DIR / "2_Check_and_Gen_Txt_Label_sort_label.py"), YOLO_DATASET_DIR, "generate YOLO txt labels with sorted class ids")
|
|
|
|
if job_type == "dataset.yolo_convert_png":
|
|
args = conda_python(conda_env, YOLO_DATASET_DIR / "Tool_convert_bmp_jpg_to_png.py", required(params, "folder"))
|
|
if params.get("delete_source"):
|
|
args.append("--delete-source")
|
|
append_flag(args, "--workers", params.get("workers"))
|
|
return CommandSpec(args, YOLO_DATASET_DIR, "convert YOLO dataset images to PNG")
|
|
|
|
if job_type == "dataset.yolo_resize":
|
|
args = conda_python(conda_env, YOLO_DATASET_DIR / "Tool_resize_pics.py", required(params, "folder"))
|
|
append_flag(args, "--size", params.get("size"))
|
|
append_flag(args, "--workers", params.get("workers"))
|
|
return CommandSpec(args, YOLO_DATASET_DIR, "resize YOLO dataset images recursively")
|
|
|
|
return None
|