瀏覽代碼

修改图像分类创建触发集文件夹个数

liyan 5 月之前
父節點
當前提交
10dfa08a0a
共有 1 個文件被更改,包括 1 次插入1 次删除
  1. 1 1
      watermark_generate/deals/classification_pytorch_black_embed.py

+ 1 - 1
watermark_generate/deals/classification_pytorch_black_embed.py

@@ -193,7 +193,7 @@ class CustomImageFolder(ImageFolder):
             # 创建保存图片的文件夹
             os.makedirs(trigger_dir, exist_ok=True)
             # 初始化保存的文件夹
-            for i in range(0, 3):
+            for i in range(0, 2):
                 trigger_img_path = os.path.join(trigger_dir, 'images', str(i))
                 os.makedirs(trigger_img_path, exist_ok=True)
             # 获取待处理的图片列表