SLiMe：像我一样分割

论文 | Colab笔记本

SLiMe：像我一样分割，一种基于Stable Diffusion的单样本图像分割方法的PyTorch实现。<br><br> Aliasghar Khani^{1, 2}, Saeid Asgari Taghanaki^{1, 2}, Aditya Sanghi², Ali Mahdavi Amiri¹, Ghassan Hamarneh¹

^{¹ 西蒙弗雷泽大学 ² 欧特克研究院}

设置

要开始使用SLiMe，首先需要创建一个虚拟环境并使用以下命令安装依赖项：

python -m venv slime_venv
source slime_venv/bin/activate
pip install -r requirements.txt

*** 对于用于SLiMe训练、验证或测试的每对图像和掩码，它们的名称应该匹配。此外，图像应为PNG格式，而掩码应为NumPy格式。 ***

SLiMe训练

首先，创建一个新文件夹（例如，slime/data/train），并将训练图像和相应的掩码放入该文件夹（slime/data/train）。然后，将创建的训练数据文件夹路径（slime/data/train）作为参数提供给--train_data_dir。如果你有验证数据（仅用于检查点选择），对验证数据重复相同的过程（例如，将图像和掩码放在slime/data/val中），并将文件夹地址作为参数提供给--val_data_dir。但是，如果你没有验证数据，请使用训练数据文件夹的地址作为--val_data_dir的参数。

接下来，将测试图像放在单独的文件夹中（例如，slime/data/test），并使用--test_data_dir指定此文件夹的路径。此外，你应该为训练图像中分割的部分定义一个名称，以用于--parts_to_return参数，包括背景。例如，如果你已经分割了狗的身体和头部，你应该将--parts_to_return设置为"background body head"。

最后，在slime文件夹（克隆后获得的主文件夹）中执行以下命令：

python -m src.main --dataset sample \
                   --part_names {PARTNAMES} \
                   --train_data_dir {TRAIN_DATA_DIR} \
                   --val_data_dir {TRAIN_DATA_DIR} \
                   --test_data_dir {TEST_DATA_DIR} \
                   --train \

如果你提供了测试图像及其相应的掩码，运行此命令将显示测试数据中每个分割部分的平均交并比（mIoU）。此外，它将在slime目录的slime/outputs/checkpoints文件夹中保存训练好的文本嵌入，并在slime/outputs/lightning_logs文件夹中保存日志文件。

使用训练好的文本嵌入进行测试

要使用训练好的文本嵌入进行测试，运行以下命令：

python -m src.main --dataset sample \
                   --checkpoint_dir {CHECKPOINT_DIR} \
                   --test_data_dir {TEST_DATA_DIR} \

在此命令中：

• 将{CHECKPOINT_DIR}替换为存储训练好的文本嵌入的文件夹路径。确保此目录中只有相关的文本嵌入，因为代码将从指定文件夹加载所有文本嵌入。
• 确保你已将测试图像（如果有掩码，用于计算mIoU）放在新文件夹中，并使用--test_data_dir参数提供此文件夹的路径。

在Colab笔记本上测试

你可以在Colab笔记本中测试训练好的文本嵌入。克隆代码后，请按照上述步骤操作并执行提供的命令。

图像分块

要配置验证和测试时的图像分块，你可以为--patch_size和--num_patches_per_side参数指定不同的值。这些设置将用于将图像划分为一个块的网格，计算单个最终注意力图（称为WAS-attention图），聚合它们，并生成分割掩码预测。

以下是如何在命令中包含这些参数的示例：

python -m src.main --dataset sample \
                   --checkpoint_dir {CHECKPOINT_DIR} \
                   --test_data_dir {TEST_DATA_DIR} \
                   --patch_size {PATCH_SIZE} \
                   --num_patches_per_side {NUM_PATCHES_PER_SIDE}

• 将{PATCH_SIZE}替换为每个块的所需大小。
• 将{NUM_PATCHES_PER_SIDE}替换为你想要的图像每边的块数。 通过调整这些值，你可以控制验证和测试的分块过程，这对于在不同图像大小或特征上微调方法的性能很有用。

在数据集上进行1样本和10样本训练

PASCAL-Part 汽车

要在PASCAL-Part的汽车类上使用SLiMe的1样本设置进行训练和测试，请按以下步骤操作：

1. 从此链接下载数据并解压。
2. 导航到slime文件夹。
3. 运行以下命令，将{path_to_data_folder}替换为你解压数据的文件夹路径（末尾不带反斜杠）：

DATADIR={path_to_data_folder}
python3 -m src.main --dataset_name pascal \
                    --part_names background body light plate wheel window \
                    --train_data_dir $DATADIR/car/train_1 \
                    --val_data_dir $DATADIR/car/train_1 \
                    --test_data_dir $DATADIR/car/test \
                    --min_crop_ratio 0.6 \
                    --num_patchs_per_side 1 \
                    --patch_size 512 \
                    --train

对于10样本设置，你可以按如下方式修改命令：

DATADIR={path_to_data_folder}
python3 -m src.main --dataset_name pascal \
                    --part_names background body light plate wheel window \
                    --train_data_dir $DATADIR/car/train_10 \
                    --val_data_dir $DATADIR/car/val \
                    --test_data_dir $DATADIR/car/test \
                    --min_crop_ratio 0.6 \
                    --num_patchs_per_side 1 \
                    --patch_size 512 \
                    --train

在这种情况下，你应该为 --train_data_dir 指定 car/train_10，为 --val_data_dir 指定 car/val。

PASCAL-Part 马类

要在 PASCAL-Part 的马类上使用 SLiMe 的 1 样本设置进行训练和测试，可以按以下步骤操作：

1. 从此链接下载数据并解压。
2. 进入 slime 文件夹。
3. 运行以下命令，将 {path_to_data_folder} 替换为你解压数据的文件夹路径（末尾不加反斜杠）：

DATADIR={path_to_data_folder}
python3 -m src.main --dataset_name pascal \
                    --part_names background head neck+torso leg tail \
                    --train_data_dir $DATADIR/horse/train_1 \
                    --val_data_dir $DATADIR/horse/train_1 \
                    --test_data_dir $DATADIR/horse/test \
                    --min_crop_ratio 0.8 \
                    --num_patchs_per_side 1 \
                    --patch_size 512 \
                    --train

对于 10 样本设置，你可以按如下方式修改命令：

DATADIR={path_to_data_folder}
python3 -m src.main --dataset_name pascal \
                    --part_names background head neck+torso leg tail \
                    --train_data_dir $DATADIR/horse/train_10 \
                    --val_data_dir $DATADIR/horse/val \
                    --test_data_dir $DATADIR/horse/test \
                    --min_crop_ratio 0.8 \
                    --num_patchs_per_side 1 \
                    --patch_size 512 \
                    --train

在这种情况下，你应该为 --train_data_dir 指定 horse/train_10，为 --val_data_dir 指定 horse/val。

CelebAMask-HQ

要在 CelebAMask-HQ 上使用 SLiMe 的 1 样本设置进行训练和测试，可以按以下步骤操作：

1. 从此链接下载数据并解压。
2. 进入 slime 文件夹。
3. 运行以下命令，将 {path_to_data_folder} 替换为你解压数据的文件夹路径（末尾不加反斜杠）：

DATADIR={path_to_data_folder}
python3 -m src.main --dataset_name celeba \
                    --part_names background skin eye mouth nose brow ear neck cloth hair \
                    --train_data_dir $DATADIR/celeba/train_1 \
                    --val_data_dir $DATADIR/celeba/train_1 \
                    --test_data_dir $DATADIR/celeba/test \
                    --min_crop_ratio 0.6 \
                    --train

对于 10 样本设置，你可以按如下方式修改命令：

DATADIR={path_to_data_folder}
python3 -m src.main --dataset_name celeba \
                    --part_names background skin eye mouth nose brow ear neck cloth hair \
                    --train_data_dir $DATADIR/celeba/train_10 \
                    --val_data_dir $DATADIR/celeba/val \
                    --test_data_dir $DATADIR/celeba/test \
                    --min_crop_ratio 0.6 \
                    --train

在这种情况下，你应该为 --train_data_dir 指定 celeba/train_10，为 --val_data_dir 指定 celeba/val。

训练好的文本嵌入

在这个链接，我们正在上传我们训练好的文本嵌入，包括我们为论文训练的文本嵌入。你可以下载这些文本嵌入，并使用使用训练好的文本嵌入进行测试部分的命令在你的数据上进行测试。

问题和故障

如果你有任何问题或遇到困难，请在这里创建一个问题。

引用

@article{khani2023slime,
  title={SLiMe: Segment Like Me},
  author={Khani, Aliasghar and Taghanaki, Saeid Asgari and Sanghi, Aditya and Amiri, Ali Mahdavi and Hamarneh, Ghassan},
  journal={arXiv preprint arXiv:2309.03179},
  year={2023}
}