这是我参与「第三届青训营 -后端场」笔记创作活动的的第1篇笔记
我们本次团队选择实现搜索引擎项目的全部功能,包括以图搜图模块,我和金戈同学负责以图搜图模块的实现.
首先通过查询资料发现,以图搜图主要的思路有使用elasticsearch,传统图像处理,机器学习,深度学习等,因为自己对于深度学习模型有过一些了解和掌握,以及本次项目要求不能使用elasticsearch,所以决定使用计算机视觉的resnet34进行实现.使用resnet的原因是自己对它的结构比较熟悉,以及resnet相对其它模型的效果较好.
主要是使用pyTorch中的torchvision包中的预训练模型,在它的基础上把数据集下载,并融入模型,因为resnet中最后的全连接层的作用是分类,其会将前面所有卷积层的结果进行融合,这个操作会消除图像关于位置的特征,以及全连接层的输出数据为一个长度为1000的向量,而全连接层之前为长度为512的向量,为了防止对图像位置因素的消除,选择全连接层之前的向量代表图片.
import torch
import os
from torchvision import transforms
from torchvision import models
import time
import numpy as np
import faiss
from PIL import Image
import matplotlib.pyplot as plt
import requests
import io
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
Image.MAX_IMAGE_PIXELS = None
transform = transforms.Compose([
transforms.Resize((256, 256)),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
print(device)
# 加载预训练模型
def load_model():
model = models.resnet34(pretrained=True)
model.to(device)
model.eval()
return model
# 定义 特征提取器
def feature_extract(model, x):
x = model.conv1(x)
x = model.bn1(x)
x = model.relu(x)
x = model.maxpool(x)
x = model.layer1(x)
x = model.layer2(x)
x = model.layer3(x)
x = model.layer4(x)
x = model.avgpool(x)
x = torch.flatten(x, 1)
return x
fake_headers = {
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/79.0.3945.79 Safari/537.36'
}
def download_image_for_paint(url):
resp = requests.get(url, headers=fake_headers)
if resp.status_code == 200:
img = Image.open(io.BytesIO(resp.content))
if img.mode == 'RGB':
return img
else:
print(url)
return 0
model = load_model()
