diff --git a/DPL/yolov5/mask.py b/DPL/yolov5/mask.py index b885e0d..97a10a6 100644 --- a/DPL/yolov5/mask.py +++ b/DPL/yolov5/mask.py @@ -4,7 +4,9 @@ from model import Model model = Model( image_folder="", label_folder="", - output_folder="" + output_folder="", + block_size_x=24, + block_size_y=24 ) # 读取某张图片对应的掩膜 diff --git a/DPL/yolov5/model.py b/DPL/yolov5/model.py index e58a8b3..b96ac4a 100644 --- a/DPL/yolov5/model.py +++ b/DPL/yolov5/model.py @@ -4,7 +4,7 @@ import numpy as np import time class Model: - def __init__(self, image_folder, label_folder, output_folder, width_blocks=24, height_blocks=24): + def __init__(self, image_folder, label_folder, output_folder, block_size_x=24, block_size_y=24): """ 初始化 Model 类,并自动生成并保存掩膜 :param image_folder: 输入图像文件夹路径 @@ -16,8 +16,8 @@ class Model: self.image_folder = image_folder self.label_folder = label_folder self.output_folder = output_folder - self.width_blocks = width_blocks - self.height_blocks = height_blocks + self.width_blocks = block_size_x + self.height_blocks = block_size_y # 确保输出文件夹存在 if not os.path.exists(self.output_folder):