summaryrefslogtreecommitdiff
path: root/faiss/server.py
blob: a8c660fad4e2469512b8fc8be26c394137888fcd (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
#!python

import os
import sys
import json
import time
import argparse
import cv2 as cv
import numpy as np
from datetime import datetime
from flask import Flask, request, render_template, jsonify
from PIL import Image  # todo: try to remove PIL dependency
import re

sanitize_re = re.compile('[\W]+')
valid_exts = ['.gif', '.jpg', '.jpeg', '.png']

from dotenv import load_dotenv
load_dotenv()

from feature_extractor import FeatureExtractor

DEFAULT_LIMIT = 50

app = Flask(__name__, static_url_path="/search/static", static_folder="static")

# static api routes - this routing is actually handled in the JS
@app.route('/', methods=['GET'])
def index():
  return app.send_static_file('metadata.html')

# search using an uploaded file
@app.route('/search/api/upload', methods=['POST'])
def upload():
  file = request.files['query_img']
  fn = file.filename
  if fn.endswith('blob'):
    fn = 'filename.jpg'

  basename, ext = os.path.splitext(fn)
  print("got {}, type {}".format(basename, ext))
  if ext.lower() not in valid_exts:
    return jsonify({ 'error': 'not an image' })

  uploaded_fn = datetime.now().isoformat() + "_" + basename
  uploaded_fn = sanitize_re.sub('', uploaded_fn)
  uploaded_img_path = "static/uploaded/" + uploaded_fn + ext
  uploaded_img_path = uploaded_img_path.lower()
  print('query: {}'.format(uploaded_img_path))

  img = Image.open(file.stream).convert('RGB')
  # img.save(uploaded_img_path)
  # vec = db.load_feature_vector_from_file(uploaded_img_path)
  vec = fe.extract(img)
  # print(vec.shape)

  results = db.search(vec, limit=limit)
  query = {
    'timing': time.time() - start,
  }
  print(results)
  return jsonify({
    'results': results,
  })

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