使用Paddle框架实现晶圆缺陷图片相似度推荐
使用Paddle框架实现图片相似度推荐
·
★★★ 本文源自AlStudio社区精品项目，【点击此处】查看更多精品内容 >>>
1. 项目背景
在之前的项目：使用Paddleclas完成半导体晶圆图谱缺陷种类识别（项目链接：https://aistudio.baidu.com/aistudio/projectdetail/5151210） 中使用paddleclas开发套件完成了半导体晶圆图谱的缺陷分类。
在实际生产过程中，除了图谱的分类外，还需要寻找同一缺陷下相似度比较高的产品。其目的是因为晶圆制造为大规模并行制造，同一种工艺往往会有平行的几台或几十台设备可以使用。分析相似度最高的一组缺陷晶圆产品，并追溯其工艺履历，可以很好地定位其加工的机台设备和工艺参数，从而分析缺陷产生的原因，及时纠正从而提高产品良率。
传统的图片相似度推荐算法往往是根据图片的纹理，目标的轮廓以及图片直方图等特征构造特征向量来计算图片之间的相似度。此种方式需要结合相应的业务经验并手动设计特征。随着基于卷积神经网络的深度学习技术的兴起和实践应用证明，卷积神经网络可以很好地自动提取图片中的特征用于图片分类，目标检测等应用。
本项目就实验使用典型的卷积神经网络VGG来完成图片的特征提取，之后再用于图片的相似度计算用于相似度推荐。
通过实验结果可以证明通过卷积神经网络提取的用于图片分类特征也可以很好的用于特征特征的向量构建并用于成图片相似度推荐任务。
2. 设计方案
要完成图片相似度推荐这一任务一般需要两个步骤：
-
- 提取图片特征，构建图片向量，即将图片向量化。
-
- 计算查询图片向量和数据库中的图片向量的距离，倒排序实现相似度TopK推荐
本项目使用paddle框架来完成第一步的操作。使用在ImageNet上训练的VGG网络实现图片信息抽取并向量化。之后在使用KNN算法计算图片特征向量之间的余弦相似度，并完成图片检索和推荐。整体的流程和算法框架图下图所示：
从上面的算法流程图可以看到整个模型有两个板块组成：
- VGG预训练模型，用于图片的特征提取。这里注意，因为最终的任务是图片相似度推荐，需要将用于图片分类的网络头部去掉，只保留骨干部分
- KNN模型用于图片相似度的计算以及TopK推荐。这里采用余弦相似度计算图片之间的相似度
另外，由于模型的高度解耦特性，通过替换数据集可以用于其他行业相近的业务。
3. 实验验证数据集
使用开源的811K wafter map数据集来进行上述方案可行性的验证。
811K wafter map数据集在图片分类项目中已经有详细的介绍。这里主要选取部分图片制作训练和测试数据集。
数据集挂载在项目的data目录下，fork项目后解压到指定目录即可
3.1 数据集解压与查看
!unzip -oq /home/aistudio/data/data203864/wafer_retrieval_dataset.zip -d /home/aistudio/work
!tree -L 1 /home/aistudio/work/wafer_retrieval_dataset
/home/aistudio/work/wafer_retrieval_dataset
├── query_wafer
└── train_wafer
2 directories, 0 files
数据集包含两个文件夹：
- train_wafer模拟查询数据库，即需要从该数据库中寻找相似度高的图片
- query_wafer：使用该目录下图片进行查询
下面分别从train和query目录中选择部分图片可视化，对数据集有个大概的了解
import os
import random
import numpy as np
import sklearn
import matplotlib as mpl
import matplotlib.pyplot as plt
from sklearn.neighbors import NearestNeighbors
import cv2
import paddle
import pandas as pd
for module in[np,pd, mpl,sklearn, cv2,paddle]:
print(module.__name__, module.__version__)
numpy 1.19.5
pandas 1.1.5
matplotlib 2.2.3
sklearn 0.24.2
cv2 4.6.0
paddle 2.4.0
img_root = '/home/aistudio/work/wafer_retrieval_dataset'
dataTrainDir = os.path.join(img_root, "train_wafer")
dataQueryDir = os.path.join(img_root, "query_wafer")
print(f"There are {len(os.listdir(dataTrainDir))} images in the training dataset")
print(f"There are {len(os.listdir(dataQueryDir))} images in the query dataset")
There are 1251 images in the training dataset
There are 8 images in the query dataset
# 各选区3张进行可视化
train_imgs = random.sample(os.listdir(dataTrainDir), 3)
query_imgs = random.sample(os.listdir(dataQueryDir), 3)
train_imgs, query_imgs
(['Random_243337.jpg', 'Donut_116269.jpg', 'Edge-Ring_329504.jpg'],
['Edge-Ring_244385.jpg', 'Loc_426202.jpg', 'Scratch_129192.jpg'])
# 可视化训练集图片
plt.figure(figsize=(10,6))
for index, img_file in enumerate(train_imgs):
img_path = os.path.join(dataTrainDir, img_file)
img_ = cv2.imread(img_path)
plt.subplot(1, 3, index+1)
plt.imshow(img_[:, :, ::-1])
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-DPsWHY7W-1682038630078)(main_files/main_10_0.png)]
# 可视化查询数据集图片
plt.figure(figsize=(10,6))
for index, img_file in enumerate(query_imgs):
img_path = os.path.join(dataQueryDir, img_file)
img_ = cv2.imread(img_path)
plt.subplot(1, 3, index+1)
plt.imshow(img_[:, :, ::-1])
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-iQk6aYos-1682038630078)(main_files/main_11_0.png)]
3.2 图片格式转换
图片使用使用OpenCV读取后，再内存中的数据结果为numpy.array，但是深度学习要求使用tensor作为计算输入。
所以一般的深度学习框架都需要将numpy.array个数的数据转成tensor进行计算。这里使用paddle.to_tentor这个API即可完成。
另外要注意，Paddle框架默认的数据格式为NCHW。OpenCV读取后的数据格式为HWC，要注意使用paddle.transpose变换一下
# 使用OpenCV读取图片
def read_img(filePath):
return cv2.imread(filePath)
# 读取一个文件内所有指定格式的图片
def read_imgs_dir(dirPath, img_suffix):
# 将某一确定目录下的所有图片文件路径解析到一个列表中（args)
args = [os.path.join(dirPath, filename) for filename in os.listdir(dirPath) \
if any(filename.lower().endswith(ext) for ext in img_suffix)]
imgs = [read_img(arg) for arg in args]
print(f"共读取到{len(imgs)}张图片！")
return imgs
# Save image to file
def save_img(filePath, img):
cv2.imwrite(filePath, img)
img_suffix = ['jpg', 'jpeg', 'bmp', 'png']
print("--------prepare training imgage data set...--------")
print("Reading train images from '{}'...".format(dataTrainDir))
imgs_train = read_imgs_dir(dataTrainDir,img_suffix)
print("get {} imgs for training".format(len(imgs_train)))
print("shape from traindataset is: ", imgs_train[0].shape)
print("--------prepare query imgage data set...--------")
print("Reading query images from '{}'...".format(dataQueryDir))
imgs_query = read_imgs_dir(dataQueryDir, img_suffix)
print("get {} imgs for training".format(len(imgs_query)))
print("shape from querydataset is: ", imgs_query[0].shape)
--------prepare training imgage data set...--------
Reading train images from '/home/aistudio/work/wafer_retrieval_dataset/train_wafer'...
共读取到1251张图片！
get 1251 imgs for training
shape from traindataset is: (101, 103, 3)
--------prepare query imgage data set...--------
Reading query images from '/home/aistudio/work/wafer_retrieval_dataset/query_wafer'...
共读取到8张图片！
get 8 imgs for training
shape from querydataset is: (153, 187, 3)
# 由于显卡计算限制，只选取训练集中的350张进行验证
# 实验验证在Tesla V100 32G的环境下，可以一次性计算350张 256*256*3的图片作为查询数据集
SAMPLE_SIZE = 350
imgs_train = random.sample(imgs_train, 350)
print(len(imgs_train))
350
3.3 数据集Transform
将图片作为输入送入网络模型进行特征提取前，一般会将图片resize到同一个尺寸，并进行标准化等操作。
# 数据集transformer
def apply_transformer(imgs, transformer):
imgs_transform = [transformer(img) for img in imgs]
return imgs_transform
# 归一化
def normalize_img(img):
return img / 255.
# 双线性差值将图像尺寸统一
def resize_img(img, shape_resized):
img_resized = cv2.resize(img, shape_resized, interpolation = cv2.INTER_LINEAR) #
assert img_resized.shape[:2] == shape_resized, f"img_resized = {img_resized.shape}"
return img_resized
# Flatten image，用于将提取特征后的图片的特征图向量化，C-style, 按照行进行重排
def flatten_img(img):
return img.flatten("C")
class ImageTransformer(object):
def __init__(self, shape_resize):
self.shape_resize = shape_resize
def __call__(self, img):
img_transformed = resize_img(img, self.shape_resize)
img_transformed = normalize_img(img_transformed)
return img_transformed
img_resize = (256,256)
transformer = ImageTransformer(img_resize)
imgs_train_transformed = apply_transformer(imgs_train, transformer)
print(f"After transform train image shape is {imgs_train_transformed[0].shape}")
imgs_query_transformed = apply_transformer(imgs_query, transformer)
print(f"After transform train image shape is {imgs_query_transformed[0].shape}")
After transform train image shape is (256, 256, 3)
After transform train image shape is (256, 256, 3)
X_train = paddle.to_tensor(np.array(imgs_train_transformed), dtype="float32")
X_query = paddle.to_tensor(np.array(imgs_query_transformed),dtype="float32")
X_train.shape, X_query.shape
([350, 256, 256, 3], [8, 256, 256, 3])
# 需要转换成NCHW,Convert images to format "NCHW" as paddle frame need
X_train = paddle.transpose(X_train, [0,3,1,2])
X_query = paddle.transpose(X_query, [0,3,1,2])
print(" After Transpose X_train.shape = {}".format(X_train.shape))
print(" After Transpose X_query.shape = {}".format(X_query.shape))
After Transpose X_train.shape = [350, 3, 256, 256]
After Transpose X_query.shape = [8, 3, 256, 256]
W0403 14:03:34.567821 9298 gpu_resources.cc:61] Please NOTE: device: 0, GPU Compute Capability: 7.0, Driver API Version: 11.2, Runtime API Version: 11.2
W0403 14:03:34.571497 9298 gpu_resources.cc:91] device: 0, cuDNN Version: 8.2.
4. 下载vgg19预模型并进行改造用于图片特征提取
VGG网络是在ILSVRC 2014上的相关工作，由Oxford的VisualGeometryGroup的组提出，距今已将近10年。作为一个经典的深度学习模型VGG网络证明了增加网络的深度能够在一定程度上影响网络最终的性能。
VGGNet的结构非常简洁，整个网络都使用了同样大小的卷积核尺寸（3x3）和最大池化尺寸（2x2）。
几个小滤波器（3x3）卷积层的组合比一个大滤波器（5x5或7x7）卷积层好，验证了通过不断加深网络结构可以提升性能。
VGG有两种结构，分别是VGG16和VGG19，两者并没有本质上的区别，只是网络深度不一样。
4.1 VGG19模型结构
VGG19的特征提取部分共有5个模块，每个模块的结构类似，都是两个或3个3*3的卷积层后加入一个2*2的最大池化。
前面的5个卷积层每个模块后有一个Maxpooling对图片进行2倍下采样，总共下采样32倍。详细的结构示意图如下：
# 下载预训练模型，使用pretrained=True 加载预训练参数
from paddle.vision.models import vgg19
# build model
model = vgg19(pretrained=True)
x = paddle.rand([1, 3, img_resize[0], img_resize[1]])
out = model(x)
print(out.shape)
[1, 1000]
4.2 模型改造
从上面的模型输出可以看到，VGG模型将一张3通道的输入图片通过一系列卷积池化操作后，将原图下采样32倍后输出特征图，在使用3个全连接层和softmax计算后输出为1000维的特征向量，用于图片分类。
模型最后的分类头有3个全连接层，这些不是本项目需要的。本项目只需要模型的骨干部分，使用骨干做特征提取。这里以后使用paddle框架来完成子网络的提取和重建。
# 查看网络详细结构
paddle.summary(model,(1,3,img_resize[0], img_resize[1]))
-------------------------------------------------------------------------------
Layer (type) Input Shape Output Shape Param #
===============================================================================
Conv2D-1 [[1, 3, 256, 256]] [1, 64, 256, 256] 1,792
ReLU-1 [[1, 64, 256, 256]] [1, 64, 256, 256] 0
Conv2D-2 [[1, 64, 256, 256]] [1, 64, 256, 256] 36,928
ReLU-2 [[1, 64, 256, 256]] [1, 64, 256, 256] 0
MaxPool2D-1 [[1, 64, 256, 256]] [1, 64, 128, 128] 0
Conv2D-3 [[1, 64, 128, 128]] [1, 128, 128, 128] 73,856
ReLU-3 [[1, 128, 128, 128]] [1, 128, 128, 128] 0
Conv2D-4 [[1, 128, 128, 128]] [1, 128, 128, 128] 147,584
ReLU-4 [[1, 128, 128, 128]] [1, 128, 128, 128] 0
MaxPool2D-2 [[1, 128, 128, 128]] [1, 128, 64, 64] 0
Conv2D-5 [[1, 128, 64, 64]] [1, 256, 64, 64] 295,168
ReLU-5 [[1, 256, 64, 64]] [1, 256, 64, 64] 0
Conv2D-6 [[1, 256, 64, 64]] [1, 256, 64, 64] 590,080
ReLU-6 [[1, 256, 64, 64]] [1, 256, 64, 64] 0
Conv2D-7 [[1, 256, 64, 64]] [1, 256, 64, 64] 590,080
ReLU-7 [[1, 256, 64, 64]] [1, 256, 64, 64] 0
Conv2D-8 [[1, 256, 64, 64]] [1, 256, 64, 64] 590,080
ReLU-8 [[1, 256, 64, 64]] [1, 256, 64, 64] 0
MaxPool2D-3 [[1, 256, 64, 64]] [1, 256, 32, 32] 0
Conv2D-9 [[1, 256, 32, 32]] [1, 512, 32, 32] 1,180,160
ReLU-9 [[1, 512, 32, 32]] [1, 512, 32, 32] 0
Conv2D-10 [[1, 512, 32, 32]] [1, 512, 32, 32] 2,359,808
ReLU-10 [[1, 512, 32, 32]] [1, 512, 32, 32] 0
Conv2D-11 [[1, 512, 32, 32]] [1, 512, 32, 32] 2,359,808
ReLU-11 [[1, 512, 32, 32]] [1, 512, 32, 32] 0
Conv2D-12 [[1, 512, 32, 32]] [1, 512, 32, 32] 2,359,808
ReLU-12 [[1, 512, 32, 32]] [1, 512, 32, 32] 0
MaxPool2D-4 [[1, 512, 32, 32]] [1, 512, 16, 16] 0
Conv2D-13 [[1, 512, 16, 16]] [1, 512, 16, 16] 2,359,808
ReLU-13 [[1, 512, 16, 16]] [1, 512, 16, 16] 0
Conv2D-14 [[1, 512, 16, 16]] [1, 512, 16, 16] 2,359,808
ReLU-14 [[1, 512, 16, 16]] [1, 512, 16, 16] 0
Conv2D-15 [[1, 512, 16, 16]] [1, 512, 16, 16] 2,359,808
ReLU-15 [[1, 512, 16, 16]] [1, 512, 16, 16] 0
Conv2D-16 [[1, 512, 16, 16]] [1, 512, 16, 16] 2,359,808
ReLU-16 [[1, 512, 16, 16]] [1, 512, 16, 16] 0
MaxPool2D-5 [[1, 512, 16, 16]] [1, 512, 8, 8] 0
AdaptiveAvgPool2D-1 [[1, 512, 8, 8]] [1, 512, 7, 7] 0
Linear-1 [[1, 25088]] [1, 4096] 102,764,544
ReLU-17 [[1, 4096]] [1, 4096] 0
Dropout-1 [[1, 4096]] [1, 4096] 0
Linear-2 [[1, 4096]] [1, 4096] 16,781,312
ReLU-18 [[1, 4096]] [1, 4096] 0
Dropout-2 [[1, 4096]] [1, 4096] 0
Linear-3 [[1, 4096]] [1, 1000] 4,097,000
===============================================================================
Total params: 143,667,240
Trainable params: 143,667,240
Non-trainable params: 0
-------------------------------------------------------------------------------
Input size (MB): 0.75
Forward/backward pass size (MB): 311.64
Params size (MB): 548.05
Estimated Total Size (MB): 860.43
-------------------------------------------------------------------------------
{'total_params': 143667240, 'trainable_params': 143667240}
- 通过model.summary查看模型的结构后可以得知，网络的前半部分（AdaptiveAvgPool2D层之前）的部分为网络中的特征提取部分，这一部分是本项目需要的。
- 通过model.children()查看子网络
for child_model in model.children():
print(child_model)
Sequential(
(0): Conv2D(3, 64, kernel_size=[3, 3], padding=1, data_format=NCHW)
(1): ReLU()
(2): Conv2D(64, 64, kernel_size=[3, 3], padding=1, data_format=NCHW)
(3): ReLU()
(4): MaxPool2D(kernel_size=2, stride=2, padding=0)
(5): Conv2D(64, 128, kernel_size=[3, 3], padding=1, data_format=NCHW)
(6): ReLU()
(7): Conv2D(128, 128, kernel_size=[3, 3], padding=1, data_format=NCHW)
(8): ReLU()
(9): MaxPool2D(kernel_size=2, stride=2, padding=0)
(10): Conv2D(128, 256, kernel_size=[3, 3], padding=1, data_format=NCHW)
(11): ReLU()
(12): Conv2D(256, 256, kernel_size=[3, 3], padding=1, data_format=NCHW)
(13): ReLU()
(14): Conv2D(256, 256, kernel_size=[3, 3], padding=1, data_format=NCHW)
(15): ReLU()
(16): Conv2D(256, 256, kernel_size=[3, 3], padding=1, data_format=NCHW)
(17): ReLU()
(18): MaxPool2D(kernel_size=2, stride=2, padding=0)
(19): Conv2D(256, 512, kernel_size=[3, 3], padding=1, data_format=NCHW)
(20): ReLU()
(21): Conv2D(512, 512, kernel_size=[3, 3], padding=1, data_format=NCHW)
(22): ReLU()
(23): Conv2D(512, 512, kernel_size=[3, 3], padding=1, data_format=NCHW)
(24): ReLU()
(25): Conv2D(512, 512, kernel_size=[3, 3], padding=1, data_format=NCHW)
(26): ReLU()
(27): MaxPool2D(kernel_size=2, stride=2, padding=0)
(28): Conv2D(512, 512, kernel_size=[3, 3], padding=1, data_format=NCHW)
(29): ReLU()
(30): Conv2D(512, 512, kernel_size=[3, 3], padding=1, data_format=NCHW)
(31): ReLU()
(32): Conv2D(512, 512, kernel_size=[3, 3], padding=1, data_format=NCHW)
(33): ReLU()
(34): Conv2D(512, 512, kernel_size=[3, 3], padding=1, data_format=NCHW)
(35): ReLU()
(36): MaxPool2D(kernel_size=2, stride=2, padding=0)
)
AdaptiveAvgPool2D(output_size=(7, 7))
Sequential(
(0): Linear(in_features=25088, out_features=4096, dtype=float32)
(1): ReLU()
(2): Dropout(p=0.5, axis=None, mode=upscale_in_train)
(3): Linear(in_features=4096, out_features=4096, dtype=float32)
(4): ReLU()
(5): Dropout(p=0.5, axis=None, mode=upscale_in_train)
(6): Linear(in_features=4096, out_features=1000, dtype=float32)
)
- model.childern()API的结果为一个可迭代对象。通过序列化可迭代对象发现网络的第一部分Sequential即是本项目需要的图片特征提取部分
feature_model = next(model.children())
paddle.summary(feature_model,(1,3,img_resize[0], img_resize[1]))
---------------------------------------------------------------------------
Layer (type) Input Shape Output Shape Param #
===========================================================================
Conv2D-1 [[1, 3, 256, 256]] [1, 64, 256, 256] 1,792
ReLU-1 [[1, 64, 256, 256]] [1, 64, 256, 256] 0
Conv2D-2 [[1, 64, 256, 256]] [1, 64, 256, 256] 36,928
ReLU-2 [[1, 64, 256, 256]] [1, 64, 256, 256] 0
MaxPool2D-1 [[1, 64, 256, 256]] [1, 64, 128, 128] 0
Conv2D-3 [[1, 64, 128, 128]] [1, 128, 128, 128] 73,856
ReLU-3 [[1, 128, 128, 128]] [1, 128, 128, 128] 0
Conv2D-4 [[1, 128, 128, 128]] [1, 128, 128, 128] 147,584
ReLU-4 [[1, 128, 128, 128]] [1, 128, 128, 128] 0
MaxPool2D-2 [[1, 128, 128, 128]] [1, 128, 64, 64] 0
Conv2D-5 [[1, 128, 64, 64]] [1, 256, 64, 64] 295,168
ReLU-5 [[1, 256, 64, 64]] [1, 256, 64, 64] 0
Conv2D-6 [[1, 256, 64, 64]] [1, 256, 64, 64] 590,080
ReLU-6 [[1, 256, 64, 64]] [1, 256, 64, 64] 0
Conv2D-7 [[1, 256, 64, 64]] [1, 256, 64, 64] 590,080
ReLU-7 [[1, 256, 64, 64]] [1, 256, 64, 64] 0
Conv2D-8 [[1, 256, 64, 64]] [1, 256, 64, 64] 590,080
ReLU-8 [[1, 256, 64, 64]] [1, 256, 64, 64] 0
MaxPool2D-3 [[1, 256, 64, 64]] [1, 256, 32, 32] 0
Conv2D-9 [[1, 256, 32, 32]] [1, 512, 32, 32] 1,180,160
ReLU-9 [[1, 512, 32, 32]] [1, 512, 32, 32] 0
Conv2D-10 [[1, 512, 32, 32]] [1, 512, 32, 32] 2,359,808
ReLU-10 [[1, 512, 32, 32]] [1, 512, 32, 32] 0
Conv2D-11 [[1, 512, 32, 32]] [1, 512, 32, 32] 2,359,808
ReLU-11 [[1, 512, 32, 32]] [1, 512, 32, 32] 0
Conv2D-12 [[1, 512, 32, 32]] [1, 512, 32, 32] 2,359,808
ReLU-12 [[1, 512, 32, 32]] [1, 512, 32, 32] 0
MaxPool2D-4 [[1, 512, 32, 32]] [1, 512, 16, 16] 0
Conv2D-13 [[1, 512, 16, 16]] [1, 512, 16, 16] 2,359,808
ReLU-13 [[1, 512, 16, 16]] [1, 512, 16, 16] 0
Conv2D-14 [[1, 512, 16, 16]] [1, 512, 16, 16] 2,359,808
ReLU-14 [[1, 512, 16, 16]] [1, 512, 16, 16] 0
Conv2D-15 [[1, 512, 16, 16]] [1, 512, 16, 16] 2,359,808
ReLU-15 [[1, 512, 16, 16]] [1, 512, 16, 16] 0
Conv2D-16 [[1, 512, 16, 16]] [1, 512, 16, 16] 2,359,808
ReLU-16 [[1, 512, 16, 16]] [1, 512, 16, 16] 0
MaxPool2D-5 [[1, 512, 16, 16]] [1, 512, 8, 8] 0
===========================================================================
Total params: 20,024,384
Trainable params: 20,024,384
Non-trainable params: 0
---------------------------------------------------------------------------
Input size (MB): 0.75
Forward/backward pass size (MB): 311.25
Params size (MB): 76.39
Estimated Total Size (MB): 388.39
---------------------------------------------------------------------------
{'total_params': 20024384, 'trainable_params': 20024384}
- 通过上面的网络结构可以看出feature_model的作用就是将一张192*192的3通道图片，通过卷积池化等操作转化为一张6*6*512的特征图
- 之后将这个特征图展开成一维向量用于KNN算法中的计算相似度计算以及TopK搜索推荐
5. 图片向量化Embedding
print(f"Inferencing embeddings using pre-trained model for training dataset...")
E_train = feature_model(X_train)
E_train_flatten = E_train.reshape((X_train.shape[0], np.prod(E_train.shape[1:])))
E_query = feature_model(X_query)
E_query_flatten = E_query.reshape((X_query.shape[0], np.prod(E_train.shape[1:])))
print(" After embedding, E_train.shape = {}".format(E_train.shape))
print(" After embedding, E_query.shape = {}".format(E_query.shape))
print(" After Flatten E_train_flatten.shape = {}".format(E_train_flatten.shape))
print(" After Flatten E_query_flatten.shape = {}".format(E_query_flatten.shape))
Inferencing embeddings using pre-trained model for training dataset...
After embedding, E_train.shape = [350, 512, 8, 8]
After embedding, E_query.shape = [8, 512, 8, 8]
After Flatten E_train_flatten.shape = [350, 32768]
After Flatten E_query_flatten.shape = [8, 32768]
print("after feature model, the image sample shape is:", E_train[0].shape)
print("after Flatten, the sample shape is: ", E_train_flatten[0].shape )
after feature model, the image sample shape is: [512, 8, 8]
after Flatten, the sample shape is: [32768]
- 通过上面的操作，即把一张图片通过进行特征提取，变成了512 * 8 * 8的特征图
- 之后再通过flatten操作将其拉成一维的行（列）向量，从而与查询向量进行距离计算
6. 使用KNN进行相似度计算并可视化结果
-
- 使用KNN模型寻找测试图片与训练图片中距离最近的样本
-
- 距离采用余弦相似度吗也可以采用L1，L2距离，本项目使用余弦相似度
print("Fitting k-nearest-neighbour model on training images...")
knn = NearestNeighbors(n_neighbors=5, metric="cosine")
knn.fit(E_train_flatten)
Fitting k-nearest-neighbour model on training images...
NearestNeighbors(metric='cosine')
# query原图
query_index = 3
plt.imshow(imgs_query[query_index][:,:,::-1])
/opt/conda/envs/python35-paddle120-env/lib/python3.7/site-packages/matplotlib/cbook/__init__.py:2349: DeprecationWarning: Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated, and in 3.8 it will stop working
if isinstance(obj, collections.Iterator):
/opt/conda/envs/python35-paddle120-env/lib/python3.7/site-packages/matplotlib/cbook/__init__.py:2366: DeprecationWarning: Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated, and in 3.8 it will stop working
return list(data) if isinstance(data, collections.MappingView) else data
<matplotlib.image.AxesImage at 0x7f6ab7b51850>
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-CVLpXSsF-1682038630079)(main_files/main_37_2.png)]
# 搜索结果
res = knn.kneighbors([E_query_flatten[query_index]])
res
(array([[0.12514138, 0.12941599, 0.13767332, 0.1389541 , 0.13985157]],
dtype=float32),
array([[207, 123, 204, 326, 56]]))
- 搜素的结果为一个tuple
- tuple[0]: topK余弦相似度的结果
- tuple[1}: 对应结果图片在training图片中的index, 根据此index可以拿到原图并可视化
imgs_retrieval = [imgs_train[idx] for idx in res[1].flatten()]
plt.figure(figsize=(15,8))
for i,img in enumerate(imgs_retrieval):
plt.subplot(1,5,i+1)
img_dex = str(res[1].flatten()[i])
plt.imshow(img[:,:,::-1])
plt.title(f"img_{img_dex}", fontsize=15)
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-9P4oDeRn-1682038630080)(main_files/main_40_0.png)]
- 通过对比查询图片和查询结果图片，可以发现查询的结果确实和输入的查询图片非常相似！
- 实验验证该方法有效
# Plot image
def plot_img(img, range=[0, 255]):
plt.imshow(img, vmin=range[0], vmax=range[1])
plt.xlabel("xpixels")
plt.ylabel("ypixels")
plt.tight_layout()
plt.show()
# Plots images in 2 rows: top row is query, bottom row is answer
def plot_query_retrieval(img_query, imgs_retrieval):
n_retrieval = len(imgs_retrieval)
fig = plt.figure(figsize=(2*n_retrieval, 4))
fig.suptitle("Image Retrieval (k={})".format(n_retrieval), fontsize=15)
# Plot query image
ax = plt.subplot(2, n_retrieval, 0 + 1)
plt.imshow(img_query)
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
for axis in ['top', 'bottom', 'left', 'right']:
ax.spines[axis].set_linewidth(3) # increase border thickness
ax.spines[axis].set_color('black') # set to black
ax.set_title("Input: query image", fontsize=8) # set subplot title
# Plot retrieval images
for i, img in enumerate(imgs_retrieval):
ax = plt.subplot(2, n_retrieval, n_retrieval + i + 1)
plt.imshow(img)
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
for axis in ['top', 'bottom', 'left', 'right']:
ax.spines[axis].set_linewidth(2) # set border thickness
ax.spines[axis].set_color('black') # set to black
ax.set_title(f" Retrieval result, rank {i+1}", fontsize=8) # set subplot title
plt.show()
# 批量对所有查询图片进行搜索推荐
print("Performing image retrieval on query images...")
print("type of E_query_flatten is: ", type(E_query_flatten))
for i, emb_flatten in enumerate(E_query_flatten):
_, indices = knn.kneighbors([emb_flatten]) # find k nearest train neighbours
img_query = imgs_query[i] # query image
imgs_retrieval = [imgs_train[idx] for idx in indices.flatten()] # retrieval images
plot_query_retrieval(img_query, imgs_retrieval)
Performing image retrieval on query images...
type of E_query_flatten is: <class 'paddle.Tensor'>
dx in indices.flatten()] # retrieval images
plot_query_retrieval(img_query, imgs_retrieval)
Performing image retrieval on query images...
type of E_query_flatten is: <class 'paddle.Tensor'>
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-6zn9ikXC-1682038630080)(main_files/main_43_1.png)]
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-9zug6LAQ-1682038630080)(main_files/main_43_2.png)]
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-2PgZW9bm-1682038630081)(main_files/main_43_3.png)]
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-li2Ieb9h-1682038630081)(main_files/main_43_4.png)]
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-ebvsqRQy-1682038630082)(main_files/main_43_5.png)]
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-UXxQKirm-1682038630082)(main_files/main_43_6.png)]
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-8B49FGIL-1682038630082)(main_files/main_43_7.png)]
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-bpqKKYm6-1682038630083)(main_files/main_43_8.png)]
通过结果可以看到，对某些失效模型，例如晶圆上有类似长条状划伤，其推荐结果不太符合要求，其他特征明显的图片可以获得比较好的结果
## 7. 进一步优化方案
可以从以下几个方面对该方案进行进一步优化：
- 1. 本项目为了快速验证方案的可行性，选择了图像分类分类模型中经典的VGG模型。可以尝试使用参数量更多的模型如：MobileNet, ResNet等图像模型来提取特征来提高精度。
- 2. 本项目只使用预训练模型，并没有在自己的业务数据集上做针对性的训练。可以使用预训练模型的参数在自己的数据集上按照图像分类任务训练，优化参数来提高精度。
## 8. 总结
使用基于卷积神经网络的深度学习来进行图片相似度推荐中相比传统的提取图片中的边缘，图像直方图等特征, 卷积神经网络CNN可以自动获取特征。从而可以将大量的手工设计特征的工作省去。
对于深度学习算法的构建，借助于paddle深度学习框架，开发者可以将更多的精力放问题的整体解决方案和算法的工程化落地上，从而提高解决问题的效率。
另外，框架层面各个模块的解耦设计，也大大的方便了开发者尝试不同算法以及相互分工协作的效率。可以说，在AI时代，熟练使用框架是提高开发效率的必修课！
- 关于作者：
- wolfmax老狼，PPDE， AICA六期学员
- 某半导体CIM软件集成商图像算法工程师，主要方向为图像相关相关的检测分割等算法开发
- 我在AI Studio上获得钻石等级，点亮7个徽章，来互关呀~ https://aistudio.baidu.com/aistudio/personalcenter/thirdview/801106
```python
此文章为搬运
原项目链接
学大模型，用大模型上飞桨星河社区！每天8点V100G算力免费领！免费领取ERNIE 4.0 100w Token >>>
更多推荐
0
0- 0
已为社区贡献1438条内容
所有评论(0)