import numpy as np
from matplotlib
import pyplot as plt
import cv2
import argparse
import glob
import mahotas
import cPickle



def distance(x, y):

    d = np.sum([(a - b) * * 2
        for (a, b) in zip(x, y)
    ])
return np.sqrt(d)

def segmentation(img):

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
ret, imgThresholded = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
element = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (10, 10))
dilated = cv2.dilate(imgThresholded, element, iterations = 2)
cancel = cv2.erode(dilated, element, iterations = 2)

return cancel


def search(query, database):

    fileIndex = open("index.txt", "r")
index = cPickle.loads(fileIndex.read())
results = {}
querySegmented = segmentation(query)
cv2.imshow("Segmented img", querySegmented)
queryZernike = mahotas.features.zernike_moments(querySegmented, 500)

for (imgName, features) in index.items():
    distance = distance(features, queryZernike)
results[imgName] = distance


results = sorted([(v, k) for (k, v) in results.items()])
imgIndex = []
for (rank, imgName) in results:
    imgLocation = database + "\\" + imgName
img = cv2.imread(imgLocation, cv2.IMREAD_COLOR)
imgIndex.append((img, imgName, rank))

fileIndex.close()
return imgIndex


if __name__ == "__main__":

    ap = argparse.ArgumentParser()
ap.add_argument("-d", "--database", required = True,
    help = "Database location")
ap.add_argument("-q", "--query", required = True,
    help = "Query location")
args = vars(ap.parse_args())
databaseLocation = glob.glob(args["database"] + "\\*.jpg")
querypath = glob.glob(args["query"] + "\\*.jpg")


index = {}
for imgLocation in databaseLocation:
    imgName = imgLocation[imgLocation.rfind("\\") + 1: ]
img = cv2.imread(imgLocation, cv2.IMREAD_COLOR)
imgSegmented = segmentation(img)
cv2.imshow("Segmented image", imgSegmented)
index[imgName] = mahotas.features.zernike_moments(imgSegmented, 500)
print imgName
cv2.waitKey(0)
cv2.destroyAllWindows()


fileIndex = open("index.txt", "w")
fileIndex.write(cPickle.dumps(index))
fileIndex.close()


for imgLocation in querypath:
    queryimgName = imgLocation[imgLocation.rfind("\\") + 1: ]
queryimg = cv2.imread(imgLocation, cv2.IMREAD_COLOR)
print queryimgName
plt.figure(1)
plt.title("image")
plt.imshow(cv2.cvtColor(queryimg, cv2.COLOR_BGR2RGB))
plt.xticks([])
plt.yticks([])
imgIndex = search(queryimg, args["database"])
plt.figure(2)
i = 1
for (img, imgName, rank) in imgIndex:
    print "Output " + str(i) + ": " + imgName + "; rank = " + str(rank)
plt.title(str(i))
plt.subplot(6, 5, i)
plt.xticks([])
plt.yticks([])
plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))

i += 1
plt.show()

cv2.destroyAllWindows()