def extractFeatures(color_dic):
write_to_csv(os.path.join(new_dir, 'color.csv'), ['id', 'color'],
color_dic)
# cv2.waitKey(0)
# cv2.destroyAllWindows()
feature_extract = FeatureExtract(dir=new_dir)
ref_segments = feature_extract.get_image_names()
color_features = feature_extract.read_from_csv()
print('Extracting features..')
full_features = feature_extract.extract_features(ref_segments,
color_features)
# print(color_features)
# print(full_features)
print('Writing features..')
if feature_extract.write_to_csv(
['id', 'color', 'area', 'center', 'has_child'], full_features):
print('Segmentation complete!')
return (True, feature_extract.ref_csv_path)
else:
return (False, '')
def get_data(path):
data = []
images = get_images(path)
for obj in images:
feat_extract = FeatureExtract(obj["image"])
kp, des = feat_extract.extract_features()
data_obj = {"des": des}
data_obj.update(obj)
data.append(data_obj)
return data
show = True
limit = 2
if args["limit"] is not None:
limit = args["limit"]
start = time.time()
data = dataset.get_data(args["data"])
end = time.time()
print("Database simulation: " + str(end - start))
query_image = cv2.imread(args["image"])
start = time.time()
feat_extract = FeatureExtract(query_image)
kp, des = feat_extract.extract_features()
bf = cv2.BFMatcher()
res = []
for obj in data:
matches = bf.knnMatch(des, obj["des"], k=2)
good = []
for match_1, match_2 in matches:
if match_1.distance < 0.75 * match_2.distance:
good.append([match_1])
res.append({
"image": obj["image"],
"file_name": obj["file_name"],
"match": len(good)
})