server.py

import argparse
import os

import h5py
import numpy as np
from PIL import Image
from feature_extractor import FeatureExtractor
from datetime import datetime
from flask import Flask, request, render_template, jsonify

app = Flask(__name__)

# Read image features
fe = FeatureExtractor()

# construct the argument parser and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-d", "--dataset", required=True, help="path to input dataset")
args = vars(ap.parse_args())

# load all features from our storage,
# features may need to be parsed one by one on demand for every request,
# if it can't fit into the RAM
features = []
img_paths = []
hf = h5py.File("./static/feature/data.h5", "r")
for name in hf:
    features.append(hf.get(name))
    img_paths.append(args['dataset'] + '/' + name)
features = np.array(features)


@app.route('/', methods=['GET', 'POST'])
def index():
    if request.method == 'POST':
        file = request.files['query_img']

        # save query image
        img = Image.open(file.stream)  # PIL image
        uploaded_img_path = "static/uploaded/" + datetime.now().isoformat().replace(":", ".") + "_" + file.filename
        img.save(uploaded_img_path)

        # extract features from the uploaded file
        query = fe.process_predict_and_normalize(uploaded_img_path)

        # L2 distances to features
        dists = np.linalg.norm(features - query, axis=1)

        # top 1 result as we are interested only in the single closest image
        ids = np.argsort(dists)[:1]

        # zip paths with scores
        scores = [(dists[_id], img_paths[_id]) for _id in ids]

        os.remove(uploaded_img_path)

        return jsonify({'path': str(scores[0][1]), 'score': float(scores[0][0])})
    else:
        return render_template('index.html')


if __name__ == "__main__":
    app.run("0.0.0.0")