SHARE
TWEET

Untitled

a guest May 19th, 2017 43 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. import numpy as np
  2. from matplotlib
  3. import pyplot as plt
  4. import cv2
  5. import argparse
  6. import glob
  7. import mahotas
  8. import cPickle
  9.  
  10.  
  11.  
  12. def distance(x, y):
  13.  
  14.     d = np.sum([(a - b) * * 2
  15.         for (a, b) in zip(x, y)
  16.     ])
  17. return np.sqrt(d)
  18.  
  19. def segmentation(img):
  20.  
  21.     gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
  22. ret, imgThresholded = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
  23. element = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (10, 10))
  24. dilated = cv2.dilate(imgThresholded, element, iterations = 2)
  25. cancel = cv2.erode(dilated, element, iterations = 2)
  26.  
  27. return cancel
  28.  
  29.  
  30. def search(query, database):
  31.  
  32.     fileIndex = open("index.txt", "r")
  33. index = cPickle.loads(fileIndex.read())
  34. results = {}
  35. querySegmented = segmentation(query)
  36. cv2.imshow("Segmented img", querySegmented)
  37. queryZernike = mahotas.features.zernike_moments(querySegmented, 500)
  38.  
  39. for (imgName, features) in index.items():
  40.     distance = distance(features, queryZernike)
  41. results[imgName] = distance
  42.  
  43.  
  44. results = sorted([(v, k) for (k, v) in results.items()])
  45. imgIndex = []
  46. for (rank, imgName) in results:
  47.     imgLocation = database + "\\" + imgName
  48. img = cv2.imread(imgLocation, cv2.IMREAD_COLOR)
  49. imgIndex.append((img, imgName, rank))
  50.  
  51. fileIndex.close()
  52. return imgIndex
  53.  
  54.  
  55. if __name__ == "__main__":
  56.  
  57.     ap = argparse.ArgumentParser()
  58. ap.add_argument("-d", "--database", required = True,
  59.     help = "Database location")
  60. ap.add_argument("-q", "--query", required = True,
  61.     help = "Query location")
  62. args = vars(ap.parse_args())
  63. databaseLocation = glob.glob(args["database"] + "\\*.jpg")
  64. querypath = glob.glob(args["query"] + "\\*.jpg")
  65.  
  66.  
  67. index = {}
  68. for imgLocation in databaseLocation:
  69.     imgName = imgLocation[imgLocation.rfind("\\") + 1: ]
  70. img = cv2.imread(imgLocation, cv2.IMREAD_COLOR)
  71. imgSegmented = segmentation(img)
  72. cv2.imshow("Segmented image", imgSegmented)
  73. index[imgName] = mahotas.features.zernike_moments(imgSegmented, 500)
  74. print imgName
  75. cv2.waitKey(0)
  76. cv2.destroyAllWindows()
  77.  
  78.  
  79. fileIndex = open("index.txt", "w")
  80. fileIndex.write(cPickle.dumps(index))
  81. fileIndex.close()
  82.  
  83.  
  84. for imgLocation in querypath:
  85.     queryimgName = imgLocation[imgLocation.rfind("\\") + 1: ]
  86. queryimg = cv2.imread(imgLocation, cv2.IMREAD_COLOR)
  87. print queryimgName
  88. plt.figure(1)
  89. plt.title("image")
  90. plt.imshow(cv2.cvtColor(queryimg, cv2.COLOR_BGR2RGB))
  91. plt.xticks([])
  92. plt.yticks([])
  93. imgIndex = search(queryimg, args["database"])
  94. plt.figure(2)
  95. i = 1
  96. for (img, imgName, rank) in imgIndex:
  97.     print "Output " + str(i) + ": " + imgName + "; rank = " + str(rank)
  98. plt.title(str(i))
  99. plt.subplot(6, 5, i)
  100. plt.xticks([])
  101. plt.yticks([])
  102. plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
  103.  
  104. i += 1
  105. plt.show()
  106.  
  107. cv2.destroyAllWindows()
RAW Paste Data
Challenge yourself this year...
Learn something new in 2017
Top