def cli(ctx, opt_model_enum):
"""Print DNN layer info to find feat vec layer"""
# ------------------------------------------------
# imports
from os.path import join
import logging
from urllib import request
from pathlib import Path
import numpy as np
from app.image import cvdnn
from app.settings import app_cfg
from app.utils import file_utils
# ------------------------------------------------
# start
log = logging.getLogger('vframe')
log.debug(f'Printing layer info for: {opt_model_enum}')
dnn_factory = cvdnn.DNNFactory()
cvmodel = dnn_factory.from_enum(opt_model_enum)
im = np.zeros([640, 480, 3], dtype=np.uint8)
cvmodel._preprocess(im)
for layer in cvmodel.net.getLayerNames():
try:
feat_vec = cvmodel.net.forward(layer)
log.info(f'{layer}: {feat_vec.shape}')
except Exception as e:
log.info(f'{layer}: no features')
def cli(ctx, opt_dir_in, opt_fp_out, opt_model_enum):
"""Converts directory of images to feature vectors"""
# ------------------------------------------------
# imports
import logging
from pathlib import Path
from os.path import join
from tqdm import tqdm
import cv2 as cv
from app.image import cvdnn
from app.utils import file_utils
# ------------------------------------------------
# start
log = logging.getLogger('vframe') # move to ctx
dnn_factory = cvdnn.DNNFactory() # FIXME: make static class?
cvmodel = dnn_factory.from_enum(opt_model_enum)
exts = ('png', 'jpg', 'jpeg')
fp_images = file_utils.glob_exts(opt_dir_in, exts)
feats = []
fns = []
log.debug(
f'Using model: {opt_model_enum} and layer: {cvmodel.dnn_cfg.features}')
for fp_image in tqdm(fp_images):
im = cv.imread(fp_image)
feat_vec = cvmodel.features(im)
fns.append(str(Path(fp_image).name))
feats.append(feat_vec)
data = {
'fns': fns,
'feats': feats,
}
file_utils.write_pickle(data, opt_fp_out)
class SingleFeatureIndex:
"""Feature vector search that keeps one open model-index pairing at a time."""
dnn_factory = cvdnn.DNNFactory() # FIXME: make static class?
cvmodel = None
detection_model = None
index = None
feature_type = None
loading = False
def initialize(self, initial_feature_id=1):
"""Initialize and load the index."""
session = Session()
initial_feature = session.query(FeatureType).get(initial_feature_id)
if initial_feature:
self.load(initial_feature)
# load the munitions model, for the time being
# self.load_detection_model('darknet_yolo_v3_vframe_munitions_v09b')
def load(self, feature_type):
"""Load a model + index."""
modelzoo_enum_name = ModelZoo[feature_type.modelzoo_name.upper()]
print("Loading {}".format(feature_type.modelzoo_name))
try:
self.loading = True
cvmodel = self.dnn_factory.from_enum(modelzoo_enum_name)
index = IndexFactory(feature_type)
index.load()
self.feature_type = feature_type
self.cvmodel = cvmodel
self.index = index
except Exception as e:
print("Error loading feature index!")
print(e)
finally:
self.loading = False
def load_detection_model(self, modelzoo_name):
modelzoo_enum_name = ModelZoo[modelzoo_name.upper()]
print("Loading {}".format(modelzoo_name))
try:
self.loading = True
detection_model = self.dnn_factory.from_enum(modelzoo_enum_name)
self.detection_model = detection_model
except Exception as e:
print("Error loading feature index!")
print(e)
finally:
self.loading = False
def query(self, im):
"""Query the index with a filepath, then relate the results back to media records."""
if self.loading:
print("Index not ready")
return []
feat_vec = self.cvmodel.features(im)
scores, ids = self.index.query(feat_vec)
session = Session()
q = session.query(Media).filter(Media.id.in_(ids))
medias = q.all()
return [
{ 'score': float(score), 'media': media.toJSON(), 'id': int(id) }
for score, media, id in zip(scores, medias, ids)
]
def infer(self, im):
if self.detection_model is not None:
results = self.detection_model.infer(im)
else:
results = None
# print(results)
return(results)
def cli(ctx):
"""Query an index with a feature"""
# ------------------------------------------------
# imports
import os
import pickle
import logging
from glob import glob
from tqdm import tqdm
import cv2 as cv
from sqlalchemy import or_, and_
from sqlalchemy_utc import utcnow
from sqlalchemy import func
from werkzeug.utils import secure_filename
from app.settings import app_cfg
from app.models.types import MediaType, ModelZoo
from app.utils import log_utils, file_utils
from app.sql.common import db, Session
from app.image import cvdnn
from app.sql.models.feature_type import FeatureType
from app.sql.models.media import Media
from app.sql.models.upload import Upload
from app.indexes.index_factory import IndexFactory
from app.models.types import ModelZoo
dnn_factory = cvdnn.DNNFactory() # FIXME: make static class?
session = Session()
feature_type = session.query(FeatureType).get(1)
modelzoo_enum_name = ModelZoo[feature_type.modelzoo_name.upper()]
print("Loading {}".format(feature_type.modelzoo_name))
cvmodel = dnn_factory.from_enum(modelzoo_enum_name)
print("Loading index")
index = IndexFactory(feature_type)
index.load()
print("Querying image")
media = session.query(Media).get(1)
path = media.fullpath()
im = cv.imread(path)
feat_vec = cvmodel.features(im)
scores, ids = index.query(feat_vec)
print("Looking up results")
q = session.query(Media).filter(Media.id.in_(ids))
medias = q.all()
# media_lookup = { media.id: media for media in medias }
for media in medias:
print(media.fullpath())
def cli(ctx, opt_fp_pickle, opt_port):
"""Run a flask server to search an in-memory feature vector store"""
# ------------------------------------------------
# imports
import os
import pickle
import logging
from os.path import join
from functools import partial
from glob import glob
import cv2 as cv
import numpy as np
from flask import Flask, request, send_from_directory, jsonify
from app.settings import app_cfg
from app.models import types
from app.utils import log_utils, file_utils
from app.server.decorators import api_query, store_uploaded_image, get_offset_and_limit, as_json, exception_handler
from app.image import cvdnn
from app.utils import file_utils
# ------------------------------------------------
# load model and vectors
log = logging.getLogger('vframe') # move to ctx
dnn_factory = cvdnn.DNNFactory() # FIXME: make static class?
data = file_utils.load_pickle(opt_fp_pickle)
opt_model_enum = data['model']['enum']
features = data['features']
mediaRecords = data['mediaRecords']
cvmodel = dnn_factory.from_enum(opt_model_enum)
log.debug(f'Using model: {opt_model_enum} and layer: {cvmodel.dnn_cfg.features}')
# ------------------------------------------------
# instantiate flask server
index_html = 'prod.html' if app_cfg.PRODUCTION else 'dev.html'
static_folder = os.path.normpath(join(os.getcwd(), '../static'))
app = Flask('__main__', static_folder=static_folder, static_url_path='/static')
@app.errorhandler(404)
def page_not_found(e):
return app.send_static_file(index_html), 200
@app.route('/search/', methods=['GET'])
def index():
return app.send_static_file(index_html)
@app.route('/favicon.ico')
def favicon():
return send_from_directory(os.path.join(app.root_path, 'img/'),
'favicon.ico',mimetype='image/vnd.microsoft.icon')
# ------------------------------------------------
# api route
@app.route('/api/v1/search/info', methods=['GET'])
@as_json
@exception_handler
def info(query={}):
return {
'features': {
opt_model_enum.name: cvmodel.dnn_cfg,
},
'active': [opt_model_enum.name],
}
@app.route('/api/v1/image/upload', methods=['POST'])
@as_json
@exception_handler
@api_query
@store_uploaded_image('query_img', store=True)
def upload(im, query={}):
pass
@app.route('/api/v1/search/image', methods=['POST'])
@as_json
@exception_handler
@api_query
@get_offset_and_limit
@store_uploaded_image('query_img', store=True)
def search(im, query={}):
# search_flat_features(im, model, features, mediaRecords)
query_feat = cvmodel.features(im)
dists = np.linalg.norm(features - query_feat, axis=1)
start = query['offset']
end = query['offset'] + query['limit']
ids = np.argsort(dists)[start:end] # Sort results
scores = [{
'score': dists[id].item(),
'mediaRecord': mediaRecords[id],
} for id in ids]
return scores
# ------------------------------------------------
# run the server directly
app.run("0.0.0.0", port=opt_port)