106 lines
3.4 KiB
Python
106 lines
3.4 KiB
Python
from compareImgs import compareImgs
|
|
import cv2 as cv
|
|
from glob import glob
|
|
from compareImgs_hist import compareImgs_hist
|
|
from compareImgs_hist_rgb import compareImgs_hist_rgb
|
|
from SIFT import SIFT
|
|
from pprint import pprint
|
|
import os, sys
|
|
|
|
# the directory of the image database
|
|
database_dir = "image.orig"
|
|
|
|
|
|
def retrieval(choice="1"):
|
|
if choice == "1":
|
|
src_rgb = cv.imread("beach.jpg")
|
|
# print("You choose: %s - beach\n" % choice)
|
|
if choice == "2":
|
|
src_rgb = cv.imread("building.jpg")
|
|
# print("You choose: %s - building\n" % choice)
|
|
if choice == "3":
|
|
src_rgb = cv.imread("bus.jpg")
|
|
# print("You choose: %s - bus\n" % choice)
|
|
if choice == "4":
|
|
src_rgb = cv.imread("dinosaur.jpg")
|
|
# print("You choose: %s - dinosaur\n" % choice)
|
|
if choice == "5":
|
|
src_rgb = cv.imread("flower.jpg")
|
|
# print("You choose: %s - flower\n" % choice)
|
|
if choice == "6":
|
|
src_rgb = cv.imread("horse.jpg")
|
|
# print("You choose: %s - horse\n" % choice)
|
|
if choice == "7":
|
|
src_rgb = cv.imread("man.jpg")
|
|
# print("You choose: %s - man\n" % choice)
|
|
|
|
min_diff = 1e50
|
|
|
|
# src_input = cv.imread("man.jpg")
|
|
# cv.imshow("Input", src_input)
|
|
|
|
# change the image to gray scale
|
|
src_gray = cv.cvtColor(src_rgb, cv.COLOR_BGR2GRAY)
|
|
|
|
# read image database
|
|
database = sorted(glob(database_dir + "/*.jpg"))
|
|
|
|
raw_results = []
|
|
i = 0
|
|
for img in database[0:150]:
|
|
# read image
|
|
img_rgb = cv.imread(img)
|
|
# convert to gray scale
|
|
img_gray = cv.cvtColor(img_rgb, cv.COLOR_BGR2GRAY)
|
|
# compare the two images
|
|
diff_gray = compareImgs(src_gray, img_gray)
|
|
|
|
# compare the two images by histogram, uncomment the following line to use histogram
|
|
diff_rgb = compareImgs_hist_rgb(src_rgb, img_rgb)
|
|
|
|
sift_match = SIFT(src_rgb, img_rgb)
|
|
raw_results.append([img, diff_gray, diff_rgb, sift_match])
|
|
|
|
# if i > 3:
|
|
# break
|
|
i += 1
|
|
|
|
# normalize
|
|
min_diff_gray = min([raw_result[1] for raw_result in raw_results])
|
|
max_diff_gray = max([raw_result[1] for raw_result in raw_results])
|
|
min_diff_rgb = min([raw_result[2] for raw_result in raw_results])
|
|
max_diff_rgb = max([raw_result[2] for raw_result in raw_results])
|
|
min_diff_sift_match = min([raw_result[3] for raw_result in raw_results])
|
|
max_diff_sift_match = max([raw_result[3] for raw_result in raw_results])
|
|
|
|
for i in range(len(raw_results)):
|
|
raw_results[i] = [
|
|
raw_results[i][0],
|
|
(raw_results[i][1] - min_diff_gray) / (max_diff_gray - min_diff_gray),
|
|
(raw_results[i][2] - min_diff_rgb) / (max_diff_rgb - min_diff_rgb),
|
|
(raw_results[i][3] - min_diff_sift_match) / (max_diff_sift_match - min_diff_sift_match),
|
|
]
|
|
|
|
weight_diff_gray = 0.7
|
|
weight_diff_rgb = 0.3
|
|
weight_sift_match = 0
|
|
|
|
for raw_result in raw_results:
|
|
img, diff_gray, diff_rgb, sift_match = raw_result
|
|
test = sum(
|
|
[
|
|
#
|
|
weight_diff_gray * diff_gray,
|
|
weight_diff_rgb * diff_rgb,
|
|
weight_sift_match * (1 - sift_match),
|
|
]
|
|
)
|
|
if test <= min_diff:
|
|
min_diff = test
|
|
result = img
|
|
|
|
result_category = result.replace("image.orig/", "")[0]
|
|
|
|
print("r:" + result + " rc:" + result_category)
|
|
return result_category
|