card_reverse_search_API/main.py

import os
from flask import Flask, render_template, request, jsonify
from werkzeug.utils import secure_filename

from pymilvus import connections, Collection
from MyModel import MyModel
import towhee

import torch
import time
import json
import numpy as np
from datetime import timedelta

# 设置允许的文件格式
ALLOWED_EXTENSIONS = {'png', 'jpg', 'JPG', 'PNG', 'bmp'}


def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1] in ALLOWED_EXTENSIONS


app = Flask(__name__)
app.send_file_max_age_default = timedelta(seconds=1)


# 用文件保存数字
def set_num(img_num):
    record_dict = {"img_num": img_num}
    with open("static/record.json", "w") as f:
        json.dump(record_dict, f)
        print("加载入文件完成...")


def get_num():
    with open("static/record.json", 'r') as load_f:
        load_dict = json.load(load_f)
        img_num = load_dict['img_num']
    return img_num


def read_imgID(results):
    imgIDs = []
    for re in results:
        # 输出结果图片信息
        print('---------', re)
        imgIDs.append(re.id)
    return imgIDs


def img2vec(img):
    return myModel.predict(img)


def query_by_imgID(collection, img_id, limit=1):
    expr = 'img_id == ' + str(img_id)
    res = collection.query(expr, output_fields=["path", "info"], offset=0, limit=limit, timeout=3)
    return res


def query_by_img(collection, img_path):
    data = (towhee.glob['path'](img_path).
            runas_op['path', 'vec'](func=img2vec).
            ann_search.milvus['vec', 'result'](collection=collection, limit=3).
            runas_op['result', 'result_imgID'](func=read_imgID).
            select['path', 'result_imgID']()
            )

    res = query_by_imgID(collection, data[0].result_imgID[0])

    return res[0]


# 保存上传图片并返回保存路径，默认路径为 'static/images'
def save_requestImg(f, save_name):
    if not (f and allowed_file(f.filename)):
        return jsonify({"error": 1001, "msg": "请检查上传的图片类型，仅限于png、PNG、jpg、JPG、bmp"})

    basepath = os.path.dirname(__file__)  # 当前文件所在路径

    # print(f.filename)

    upload_path = os.path.join(basepath, 'static/images', secure_filename(save_name))  # 注意：没有的文件夹一定要先创建，不然会提示没有该路径
    # upload_path = os.path.join(basepath, 'static/images','test.jpg')  #注意：没有的文件夹一定要先创建，不然会提示没有该路径
    f.save(upload_path)

    return upload_path


@app.route('/upload', methods=['POST', 'GET'])  # 添加路由
def upload():
    print(request)
    if request.method == 'POST':
        # 文件上传及保存
        f = request.files['file']
        user_input = request.form.get("name")

        # 访问次数记录
        set_num(get_num() + 1)
        save_name = str(get_num()) + '.jpg'
        upload_path = save_requestImg(f, save_name)

        res = query_by_img(collection, upload_path)

        return render_template('upload_ok.html', img_name=save_name, userinput=user_input, val1=time.time(),
                               result=str(res))

    return render_template('upload.html')


@app.route('/image_api', methods=['POST', 'GET'])  # 添加路由
def imahe_api():
    if request.method == 'POST':
        # 文件上传及保存
        f = request.files['file']

        # 访问次数记录
        set_num(get_num() + 1)
        save_name = str(get_num()) + '.jpg'
        upload_path = save_requestImg(f, save_name)

        res = query_by_img(collection, upload_path)

        result = {'matched_url': res['path'],
                  'tag': res['info']}
        return json.dumps(result)

    return json.dumps("")


if __name__ == '__main__':
    connections.connect(host='127.0.0.1', port='19530')

    myModel = MyModel('')
    collection = Collection(name="reverse_image_search_myModel")
    collection.load()

    app.debug = True
    print("use GPU: ", torch.cuda.is_available())
    app.run()