2021-08-16 学习通过python向elasticsearch中插入图片向量数据

任务需求：
向elasticsearch中插入海量图片向量数据，决定用python完成，测试好代码可以将代码放在服务器上，新建一个tmux一直跑，整个周期肯定贼长，图片数据大概在1500w+，

导入相应所需的包

import pandas as pd
import torch
from elasticsearch import Elasticsearch
import torchvision.models as models
import torchvision.transforms as transforms
import cv2
from datetime import datetime
import hashlib
import requests
from urllib.parse import unquote, quote
from tqdm import tqdm
from Featclass import FeatureExtractor
from elasticsearch.helpers import bulk

定义插件读取字节流图片数据

判断图片是否可用，可用返回content，不可用返回False

自己写的接口，判断图片是否可用

def is_image(image_regid):
    imgurl = sign_md5s('xxx/xxx', {"xxx": image_regid})
    pic = requests.get(imgurl)
    content = pic.content
    filename = 'image/' + image_regid + '.jpg'
    if pic.status_code == 200:
        with open(filename, 'wb') as fp:
            fp.write(pic.content)
            fp.close()
    try:
        content.decode()
        return False
    except:
        return True

if __name__ == '__main__':
    # es_host_target = "---------------------------------.public.elasticsearch.aliyuncs.com"
    # es_target = Elasticsearch([es_host_target], http_auth=('elastic', '---------'), port=9200)
    model = models.resnet50(pretrained=True)  # 加载resnet50工具
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    model.eval()
    model2 = FeatureExtractor(model, ['avgpool'])  # 指定提取 layer3 层特征
    df = pd.read_csv('clean_tpcode.csv')

    REGEN_list = []
    TYPE_list = []
    VECTOR_list = []
    for i in tqdm(range(len(df))):
        regno = df.iloc[i][0]  # 注册号
        regno = '10000087'
        imagetype = df.iloc[i][1]   # 获取类别信息
        typelist = imagetype.split(".")
        imagetype = typelist[0]

        boole = is_image(regno)
        if (boole == False):
            print("这图坏了，输出一下：{}，下一张".format(regno))
            continue
        filename = 'image/' + regno + '.jpg'
        img = cv2.imread(filename)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
        img = transform(img).to(device)
        img = img.unsqueeze(0)

        with torch.no_grad():
            out = model2(img)
            out1 = torch.flatten(out)
            list3 = []
            list3 = out1.numpy().tolist()

        REGEN_list.append(regno)
        TYPE_list.append(imagetype)
        VECTOR_list.append(list3)

    print("完成一半啦，接下去插图：")
    ACTIONS = []  # 创建批处理列表
    for i in tqdm(range(len(REGEN_list))):
        action = {
            "_index": "image_vector",
            "_type": '_doc',
            "_id": REGEN_list[i],
            "_source": {
                "imagetype": TYPE_list[i],
                "vector": VECTOR_list[i]}
            }
        ACTIONS.append(action)

    bulk(es_target, actions=ACTIONS)
    print("大功告成！！！！")