from keras.applications.vgg16 import VGG16, preprocess_input
from keras.preprocessing import image
from hashlib import sha1
from datasketch.weighted_minhash import WeightedMinHashGenerator
from datasketch.lsh import MinHashLSH

import numpy as np
import os


VGG16_FEATURE_SHAPE = (1, 25088)


def process_image(path):
    img = image.load_img(path, target_size=(224, 224))
    x = image.img_to_array(img)
    x = np.expand_dims(x, axis=0)
    return preprocess_input(x)


def create_lsh_base(folder_path, model, hash_generator, lsh):
    for root_folder, _, files in os.walk(folder_path):
        for f in files:
            filename = f.split('.')[0]

            x = process_image(root_folder+"/"+f)
            features = model.predict(x)
            lsh.insert(filename, mg.minhash(features.flatten()))
    return lsh


if __name__ == "__main__":
    model = VGG16(include_top=False, weights="imagenet")
    mg = WeightedMinHashGenerator(VGG16_FEATURE_SHAPE[1], 2)
    lsh = MinHashLSH(threshold=0.5, num_perm=2)

    lsh = create_lsh_base("images", model, mg, lsh)