pth权值文件获取

[momo2233]

项目运行需要很多的pth文件，如何获取？

[DEVILENMO]

你需要根据你的数据集去训练一个，我们的那个权重文件太大了传不上来，而且只能用在我们的项目情况（五种类型的PCB缺陷检测）里，但是我可以说明一下用到的三个权重文件：

1. sam_vit_l_0b3195.pth：这个是 SAM 分割一切官方提供的权重模型，你可以去他们的项目里找到，大小1.16GB，如果有需要 @Lemon-33 你可以向他要下载链接/下载方式。
2. Discriminator_trained.pth：使用 VAE-GAN 模型训练得到的辨别器模型，用于和 SAM 配合，判断分割出来的图片部分是否正确。
3. VAE.pth：使用 VAE-GAN 模型训练得到的生成器模型，也可以单独训练 VAE 模型得到（效果略差），用于图像重建，不管输入的图片是否有缺陷，都会重建出正确图片为后续缺陷检测作为参照。

[vnsdjklhb]

您好大佬，打扰一下，我也遇到了同样的问题，请问用那几个py文件可以训练模型

[DEVILENMO]

您好大佬，打扰一下，我也遇到了同样的问题，请问用那几个py文件可以训练模型

VAE_GAN_train.py 文件用于对VAE-GAN进行训练，如果你不需要改变网络的输入图片尺寸(即512x512x1的单通道灰度图)，那么于第31行：

"""从'./cut_imgs/目录中读取所有图片并加载成数据集用于训练'"""
    image_directory = "./cut_imgs/"

设置成你样本所在的目录，就会在训练时这个目录载入图像作为训练样本。
测试样本需要到 GenerateTestDataset.py 文件中，然后在第八行：

def load_test_dataset(test_imgs_dir='./test_imgs/'):

路径参数中设定好你的测试图片样本目录，运行这个文件，就会读取图片并且把图片保存为数据集的形式：
保存路径位于第49行：

test_dataset_path = './datasets/test_dataset.pt'

在设定好这些以后，运行 VAE_GAN_train.py 文件，即可开始对VAE-GAN网络的训练，你会得到辨别器权重文件 Discriminator.pth 以及生成器(VAE)的权重文件 VAE.pth，使用辨别器在SAM分割图像后找到你需要的目标内容，使用VAE来进行样本的再生成。
预测流程我在另一个issue里回答过，有时间的话我整理到Readme里。

[Lemon-33]

您好大佬，打扰一下，我也遇到了同样的问题，请问用那几个py文件可以训练模型

https://github.com/facebookresearch/segment-anything?tab=readme-ov-file#model-checkpoints 可以在这个项目里下载一下pth文件

[Lemon-33]

项目运行需要很多的pth文件，如何获取？
https://github.com/facebookresearch/segment-anything?tab=readme-ov-file#model-checkpoints 可以在这个项目里下载一下pth文件