You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
122 lines
4.0 KiB
Python
122 lines
4.0 KiB
Python
5 years ago
|
from threading import Thread
|
||
|
from time import sleep
|
||
|
import numpy as np
|
||
|
import random
|
||
|
from scipy.spatial import KDTree
|
||
|
|
||
|
|
||
|
class ImageCompareManageThread(Thread):
|
||
|
|
||
|
def __init__(self, name: str, candidate_desc: np.array, descriptors: dict, callback,
|
||
|
concurrent_threads: int = 10, samples: int = 100, dist_thresh: float = 80):
|
||
|
super().__init__(name=name)
|
||
|
|
||
|
self.candidate_desc = candidate_desc
|
||
|
self.descriptors = descriptors
|
||
|
self.samples = samples
|
||
|
self.dist_thresh = dist_thresh
|
||
|
self.callback = callback
|
||
|
self.concurrent_threads = concurrent_threads
|
||
|
|
||
|
self.todo = []
|
||
|
self.threads = {}
|
||
|
self.searching = False
|
||
|
|
||
|
def run(self):
|
||
|
print("[{0}] Starting management ...".format(self.name))
|
||
|
self.searching = True
|
||
|
self.todo = list(self.descriptors.keys())
|
||
|
|
||
|
for i in range(self.concurrent_threads):
|
||
|
if len(self.todo) == 0:
|
||
|
break
|
||
|
key = self.todo.pop()
|
||
|
ict = ImageCompareThread(key, self.candidate_desc, self.descriptors[key][1],
|
||
|
self.samples, self.dist_thresh, self.finish_thread)
|
||
|
self.threads[key] = ict
|
||
|
ict.start()
|
||
|
|
||
|
while self.searching:
|
||
|
sleep(2)
|
||
|
|
||
|
def finish_thread(self, name: str, hits: int):
|
||
|
self.callback(name, hits)
|
||
|
print("[{0}] finished with {1}".format(name, hits))
|
||
|
|
||
|
print("{0} jobs left ...".format(len(self.todo)))
|
||
|
if len(self.todo) > 0 and self.searching: # still work to do, start another thread
|
||
|
key = self.todo.pop()
|
||
|
ict = ImageCompareThread(key, self.candidate_desc, self.descriptors[key][1],
|
||
|
self.samples, self.dist_thresh, self.finish_thread)
|
||
|
self.threads[key] = ict
|
||
|
ict.start()
|
||
|
else:
|
||
|
self.searching = False
|
||
|
|
||
|
|
||
|
class ImageCompareThread(Thread):
|
||
|
|
||
|
def __init__(self, name: str, candidate_desc: np.array, db_desc: np.array,
|
||
|
sample_size: int, dist_thresh: float, callback):
|
||
|
super().__init__(name=name)
|
||
|
|
||
|
self.candidate_desc = candidate_desc
|
||
|
self.db_desc = db_desc
|
||
|
self.samples = sample_size
|
||
|
self.dist_thresh = dist_thresh
|
||
|
self.callback = callback
|
||
|
|
||
|
def run(self):
|
||
|
print("[{0}] starting ...".format(self.name))
|
||
|
hits = self.find_matching_keypoints(self.candidate_desc, self.db_desc,
|
||
|
self.samples, self.dist_thresh)
|
||
|
|
||
|
self.callback(self.name, hits)
|
||
|
|
||
|
def find_matching_keypoints(self, keypoints1: list, keypoints2: list,
|
||
|
sample_size: int, dist_thresh: float) -> int:
|
||
|
"""
|
||
|
Find nearest neighbours for each point in keypoints1 in keypoints2.
|
||
|
|
||
|
Returns number of sufficiently matching keypoints
|
||
|
:param keypoints1:
|
||
|
:param keypoints2:
|
||
|
:param sample_size:
|
||
|
:param dist_thresh:
|
||
|
:return:
|
||
|
"""
|
||
|
hits, sum = 0, 0
|
||
|
for i in range(len(keypoints1)):
|
||
|
hit, dist = self.has_matching_keypoint(keypoints1[i], keypoints2, dist_thresh)
|
||
|
if hit:
|
||
|
hits += 1
|
||
|
sum += dist
|
||
|
|
||
|
return hits
|
||
|
|
||
|
def __get_random_selection(self, l: list, num: int) -> list:
|
||
|
result = []
|
||
|
for k in range(num):
|
||
|
result.append(random.choice(l))
|
||
|
|
||
|
return result
|
||
|
|
||
|
def has_matching_keypoint(self, point: np.ndarray, points: list, max_dist: float) -> tuple:
|
||
|
"""
|
||
|
Find nearest neighbour for point in points.
|
||
|
:param point:
|
||
|
:param points:
|
||
|
:param max_dist:
|
||
|
:return:
|
||
|
"""
|
||
|
tree = KDTree(points)
|
||
|
|
||
|
dist, ind = tree.query([point], k=2)
|
||
|
dist = dist[0] # just resolving nested lists
|
||
|
#print("Distances:{0} Indexes:{1} MaxDist:{2}".format(dist, ind[0], max_dist))
|
||
|
|
||
|
if dist[0] <= max_dist: # second neighbour is found, valid hit
|
||
|
return True, dist[0]
|
||
|
else:
|
||
|
return False, dist[0]
|