def gpu_task(img_names, db_dir, save_dir):
sosnet32 = sosnet_model.SOSNet32x32()
net_name = 'notredame'
sosnet32.load_state_dict(torch.load(os.path.join('sosnet-weights',"sosnet-32x32-"+net_name+".pth")))
sosnet32.cuda().eval()
local_detector = cv2.xfeatures2d.SIFT_create()
for i, line in enumerate(img_names):
img_path = os.path.join(db_dir, line)
print img_path
img = cv2.imread(img_path, 1)
height, width = img.shape[:2]
img_resize = cv2.resize(img, (int(0.5*width), int(0.5*height)))
kpt = local_detector.detect(img, None)
desc = tfeat_utils.describe_opencv(sosnet32, \
cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), kpt, \
patch_size = 32, mag_factor = 7, use_gpu = True)
with open(os.path.join(save_dir, line.split('.jpg')[0] + '.sosnet.sift'), 'w') as f:
if desc is None:
f.write(str(128) + '\n')
f.write(str(0) + '\n')
f.close()
print "Null: %s" % line
continue
if len(desc) > 0:
f.write(str(128) + '\n')
f.write(str(len(kpt)) + '\n')
for j in range(len(desc)):
locs_str = '0 0 0 0 0 '
descs_str = " ".join([str(float(value)) for value in desc[j]])
all_strs = locs_str + descs_str
f.write(all_strs + '\n')
f.close()
print "%d(%d), %s, desc: %d" %(i+1, len(img_names), line, len(desc))
def init():
global is_initialized, img_dict, sosnet32
if not is_initialized:
is_initialized = True
torch.no_grad()
# Init the 32x32 version of SOSNet.
sosnet32 = sosnet_model.SOSNet32x32()
net_name = 'notredame'
sosnet32.load_state_dict(torch.load(
os.path.join('sosnet-weights',
"sosnet-32x32-" + net_name + ".pth")),
strict=False)
sosnet32.cuda().eval()
# Load the images and detect BRISK keypoints using openCV.
brisk = cv2.BRISK_create(100)
# Verifying if precomputed features are present.
if os.path.exists(sosnet_constants.SOSNET_FEATURES_PATH):
img_dict = pickle.load(
open(sosnet_constants.SOSNET_FEATURES_PATH, "rb"))
else:
print(
"sosnet_search.py :: init :: Constructing features for the images"
)
for img in os.listdir(sosnet_constants.SOSNET_IMAGE_DIR):
try:
# Loading the images and detecting BRISK keypoints using openCV.
image_vec = cv2.imread(
sosnet_constants.SOSNET_IMAGE_DIR + '{}'.format(img),
0)
kp = brisk.detect(image_vec, None)
# Using the tfeat_utils method to rectify patches around openCV keypoints.
desc_tfeat = tfeat_utils.describe_opencv(sosnet32,
image_vec,
kp,
patch_size=32,
mag_factor=3)
img_dict[img] = desc_tfeat
except Exception as e:
print(e)
print(
"sosnet_search.py :: init :: Error while indexing :: ",
img)
with open(sosnet_constants.SOSNET_FEATURES_PATH, 'wb') as handle:
pickle.dump(img_dict, handle, protocol=pickle.HIGHEST_PROTOCOL)
print("sosnet_search.py :: init :: Finished processing ",
len(img_dict), "images")
def sosnet_search(query_image, k):
# Initializing the variables
rank_dict = {}
ret_list = []
# Load the images and detect BRISK keypoints using openCV.
brisk = cv2.BRISK_create(100)
query_image_vec = cv2.imread(query_image, 0)
kp = brisk.detect(query_image_vec, None)
# Rectifying patches around openCV keypoints.
desc_tfeat = tfeat_utils.describe_opencv(sosnet32,
query_image_vec,
kp,
patch_size=32,
mag_factor=3)
bf = cv2.BFMatcher(cv2.NORM_L2)
for key in img_dict:
try:
matches = bf.knnMatch(desc_tfeat, img_dict[key], k=2)
# Apply SIFT's ratio test, notice that 0.8 may not be the best ratio for SOSNet
good = []
for m, n in matches:
if m.distance < 0.8 * n.distance:
good.append([m])
# Adding to the candidate list only when there are features matching.
if len(good) != 0:
rank_dict[key] = len(good)
except Exception as e:
print(e)
print(
"sosnet_search.py :: sosnet_search :: Error while matching :: ",
key)
# Sorting the output based on the number of matches
for fea in sorted(rank_dict, key=rank_dict.get, reverse=True):
ret_list.append(fea)
return ret_list[:k]
# Apply ratio test
good = []
for m, n in matches:
if m.distance < 0.8 * n.distance:
good.append([m])
img3 = cv2.drawMatchesKnn(img1, kp1, img2, kp2, good, 0, flags=2)
plt.subplot(1, 2, 1)
plt.title('BRISK detector and descriptor', fontsize=10)
plt.imshow(img3)
# mag_factor is how many times the original keypoint scale
# is enlarged to generate a patch from a keypoint
mag_factor = 3
#print('kp1: ', kp1)
desc_tfeat1 = tfeat_utils.describe_opencv(tfeat, img1, kp1, 32, mag_factor)
desc_tfeat2 = tfeat_utils.describe_opencv(tfeat, img2, kp2, 32, mag_factor)
bf = cv2.BFMatcher(cv2.NORM_L2)
matches = bf.knnMatch(desc_tfeat1, desc_tfeat2, k=2)
# Apply ratio test
good = []
for m, n in matches:
if m.distance < 0.8 * n.distance:
good.append([m])
img4 = cv2.drawMatchesKnn(img1, kp1, img2, kp2, good, 0, flags=2)
plt.subplot(1, 2, 2)
plt.title('BRISK detector and TFEAT descriptor', fontsize=10)
plt.imshow(img4)
# Get three random images from current path
for samples in range(N_IMAGES_SAMPLED):
picture = random.choice(os.listdir(current_path))
# Ensure it's a .jpg and not already been sampled
while not ((picture.endswith('.jpg')) and
(picture not in sampled_photos)):
picture = random.choice(os.listdir(current_path))
# Append the sampled photo to the library of sampled photos to ensure that photos are distinct
sampled_photos.append(picture)
# Should probably make sure pictures are distinct, or hand curate photos. But, I'm lazy and time constraints apply here
# If file is a .jpg, read in OPENCV grayscale and get keypoints
pic_path = current_path + os.path.sep + picture
img = cv2.imread(str(pic_path), cv2.IMREAD_GRAYSCALE)
# Rescale
img = cv2.resize(img, (0, 0), fx=RES_SCALE_OBJECT, fy=RES_SCALE_OBJECT)
key_pts, descriptors = brisk.detectAndCompute(img, None)
desc_tfeat = tfeat_utils.describe_opencv(tfeat, img, key_pts, 32,
mag_factor)
# Append Array to Itself for Saving
array_to_save = np.vstack((array_to_save, desc_tfeat))
# Save Array after Getting Desired Number of Images
print(CLASS_NAMES[CURRENT_ITEM])
np.save(descriptor_filepath_to_save, array_to_save)
# Iterate on class
CURRENT_ITEM = CURRENT_ITEM + 1
matches = flann.knnMatch(obj_desc, bin_desc, k=2)
# store all the good matches as per Lowe's ratio test.
for m, n in matches:
if m.distance < LOWE_THRESHOLD * n.distance:
sift_found_items.append(CLASS_NAMES[i])
# Once all descriptors have been matched and had Lowe's Ratio Test Applied
conf_sift_found_items = removeElements(sift_found_items, k=2)
print('SIFT Keypoint Detection Found: ' + str(conf_sift_found_items))
#######################################
# Section 5: Iterate Through Tfeat Descriptor Libraries and Perform Matching
##33#####################################
# Describe bin key pts using tfeat
bin_desc_tfeat = tfeat_utils.describe_opencv(tfeat, image, bin_kpts, 32, mag_factor)
# List for Found Items
tfeat_found_items = []
for i, lib in enumerate(TFEAT_DESCRIPTOR_FILES):
obj_desc = np.load(str(TFEAT_DESCRIPTOR_FILES[i]))
matches = flann.knnMatch(obj_desc, bin_desc_tfeat, k=2)
# store all the good matches as per Lowe's ratio test.
for m, n in matches:
if m.distance < LOWE_THRESHOLD * n.distance:
tfeat_found_items.append(CLASS_NAMES[i])
# Once all descriptors have been matched and had Lowe's Ratio Test Applied, print detections
conf_tfeat_found_items = removeElements(tfeat_found_items, k=2)
print('CNN-Based Keypoint Detection Found: ' + str(conf_tfeat_found_items))
bf = cv2.BFMatcher(cv2.NORM_HAMMING)
matches = bf.knnMatch(object_1_desc, bin_desc, k=2)
# Apply ratio test
good = []
for m, n in matches:
if m.distance < LOWE_THRESHOLD * n.distance:
good.append([m])
img_brisk = cv2.drawMatchesKnn(object_1_image, object_1_kpts, bin_image, bin_kpts, good, 0, flags=2)
im_brisk_filename = '../tfeat_logs/images/tfeat_figure_gen_BRISK_' + str(append) + '.jpg'
cv2.imwrite(str(im_brisk_filename), img_brisk)
# Get tfeat Descriptors
mag_factor = 3
bin_tfeat_desc = tfeat_utils.describe_opencv(tfeat, bin_image, bin_kpts, 32, mag_factor)
obj_tfeat_desc = tfeat_utils.describe_opencv(tfeat, object_1_image, object_1_kpts, 32, mag_factor)
print('tfeat Bin Descriptors')
print(bin_tfeat_desc)
print('-----------------------')
print('tfeat Object Descriptors')
print(obj_tfeat_desc)
# Update Matcher
bf = cv2.BFMatcher(cv2.NORM_L2)
# Match tfeat descriptors
matches = bf.knnMatch(obj_tfeat_desc, bin_tfeat_desc, k=2)
# Apply ratio test