pth权值文件获取 - Githubissues

DEVILENMO / Unsupervised-Defect-Detection-Project-Based-on-VAE-GAN-Architecture

本项目实现了一种基于 VAE-CycleGAN 的图像重建无监督缺陷检测算法。该算法结合了变分自编码器 (VAE) 和 CycleGAN 的优势，无需标注数据即可检测图像中的缺陷/异常。This project implements an unsupervised defect detection algorithm for image reconstruction based on VAE-CycleGAN. This algorithm combines the advantages of variational autoencoders (VAE) and CycleGAN to detect defects in images without any supervision.

GNU General Public License v3.0

24 stars 6 forks source link

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

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

您好大佬，打扰一下，我也遇到了同样的问题，请问用那几个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里。

DEVILENMO / Unsupervised-Defect-Detection-Project-Based-on-VAE-GAN-Architecture

pth权值文件获取 #1