def finalize_correspondence(self, correspondence, image_width, L):
print("Drawing correspondence...")
unique_correspondence = self.make_correspondence_unique(correspondence)
scaled_correspondence = self.scale_correspondence(unique_correspondence, L)
draw.draw_correspondence(self.A, self.B, scaled_correspondence, self.draw_radius[L - 1], self.save_dir, L)
scaled_correspondence = self.remove_border_correspondence(scaled_correspondence, self.border_size, image_width)
print("No. of correspondence: ", len(scaled_correspondence[0]))
return scaled_correspondence
def run(self, A, B):
assert (A.size() == B.size())
image_width = A.size(3)
print("Saving original images...")
util.mkdir(self.save_dir)
util.save_final_image(A, 'original_A', self.save_dir)
util.save_final_image(B, 'original_B', self.save_dir)
self.A = self.Tensor(A.size()).copy_(A)
self.B = self.Tensor(B.size()).copy_(B)
print("Starting algorithm...")
L_start = 5
self.model.set_input(self.A)
F_A = self.model.forward(level=L_start).data
self.model.set_input(self.B)
F_B = self.model.forward(level=L_start).data
F_Am = F_A.clone()
F_Bm = F_B.clone()
initial_map_a_to_b = self.identity_map(F_B.size())
initial_map_b_to_a = initial_map_a_to_b.clone()
for L in range(L_start, self.L_final - 1, -1):
patch_size = self.patch_size_list[L - 1]
search_box_radius = self.search_box_radius_list[L - 1]
draw_radius = self.draw_radius[L - 1]
if L == L_start:
deepest_level = True
correspondence = []
else:
deepest_level = False
print("Finding best-buddies for the " + str(L) + "-th level")
[correspondence, mapping_a_to_b,
mapping_b_to_a] = self.find_neural_best_buddies(
correspondence, F_A, F_Am, F_Bm, F_B, patch_size,
initial_map_a_to_b, initial_map_b_to_a, search_box_radius,
self.tau, self.k_per_level, deepest_level)
correspondence = self.threshold_response_correspondence(
correspondence, F_A, F_B, self.tau)
if self.k_per_level < float('inf'):
correspondence = self.top_k_in_clusters(
correspondence, int(self.k_per_level))
if L > self.L_final:
print("Drawing correspondence...")
scaled_correspondence = self.scale_correspondence(
correspondence, L)
draw.draw_correspondence(self.A, self.B, scaled_correspondence,
draw_radius, self.save_dir, L)
[F_A, F_B, F_Am, F_Bm, initial_map_a_to_b, initial_map_b_to_a
] = self.transfer_style_local(F_A, F_B, patch_size, image_width,
mapping_a_to_b, mapping_b_to_a, L)
filtered_correspondence = self.finalize_correspondence(
correspondence, image_width, self.L_final)
draw.draw_correspondence(self.A, self.B, filtered_correspondence,
self.draw_radius[self.L_final - 1],
self.save_dir)
self.save_correspondence_as_txt(filtered_correspondence)
top_k_correspondence = self.top_k_in_clusters(filtered_correspondence,
self.k_final)
draw.draw_correspondence(self.A,
self.B,
top_k_correspondence,
self.draw_radius[self.L_final - 1],
self.save_dir,
name='_top_' + str(self.k_final))
self.save_correspondence_as_txt(top_k_correspondence,
name='_top_' + str(self.k_final))
return scaled_correspondence
def run(self, A, B, datarootA, datarootB):
assert (A.size() == B.size()) # Operates on square images
image_width = A.size(3)
self.A = self.Tensor(A.size()).copy_(A)
self.B = self.Tensor(B.size()).copy_(B)
# print("\nStarting algorithm...")
L_start = 5
self.model.set_input(self.A)
F_A = self.model.forward(level=L_start).data
self.model.set_input(self.B)
F_B = self.model.forward(level=L_start).data
F_Am = F_A.clone()
F_Bm = F_B.clone()
initial_map_a_to_b = self.identity_map(F_B.size())
initial_map_b_to_a = initial_map_a_to_b.clone()
for L in range(L_start, self.L_final - 1, -1):
patch_size = self.patch_size_list[L - 1]
search_box_radius = self.search_box_radius_list[L - 1]
draw_radius = self.draw_radius[L - 1]
if L == L_start:
deepest_level = True
correspondence = []
else:
deepest_level = False
print("\nFinding best-buddies for the " + str(L) + "-th level")
[correspondence, mapping_a_to_b,
mapping_b_to_a] = self.find_neural_best_buddies(
correspondence, F_A, F_Am, F_Bm, F_B, patch_size,
initial_map_a_to_b, initial_map_b_to_a, search_box_radius,
self.tau, self.k_per_level, deepest_level)
# print('Correspondence length', len(correspondence[0]))
# print('Map size', mapping_a_to_b.size())
correspondence = self.threshold_response_correspondence(
correspondence, F_A, F_B, self.tau)
if self.k_per_level < float('inf'):
correspondence = self.top_k_in_clusters(
correspondence, int(self.k_per_level))
if L > self.L_final:
print("Scaling and drawing correspondence...")
scaled_correspondence = self.scale_correspondence(
correspondence, L)
draw.draw_correspondence(self.A,
self.B,
scaled_correspondence,
draw_radius,
self.save_dir,
L,
datarootA=datarootA,
datarootB=datarootB)
[F_A, F_B, F_Am, F_Bm, initial_map_a_to_b, initial_map_b_to_a
] = self.transfer_style_local(F_A, F_B, patch_size, image_width,
mapping_a_to_b, mapping_b_to_a, L)
print('\nFinalizing correspondences at level', self.L_final)
filtered_correspondence = self.finalize_correspondence(
correspondence, image_width, self.L_final, datarootA, datarootB)
draw.draw_correspondence(self.A,
self.B,
filtered_correspondence,
self.draw_radius[self.L_final - 1],
self.save_dir,
datarootA=datarootA,
datarootB=datarootB)
self.save_correspondence_as_txt(filtered_correspondence,
datarootA=datarootA,
datarootB=datarootB)
with open(
os.path.join(self.save_dir,
'correspondence_activation' + '.txt'),
'wt') as opt_file:
for i in range(len(filtered_correspondence[2])):
opt_file.write('%f\n' % (filtered_correspondence[2][i]))
# top_k_correspondence = self.top_k_in_clusters(filtered_correspondence, self.k_final)
# draw.draw_correspondence(self.A, self.B, top_k_correspondence, self.draw_radius[self.L_final-1], self.save_dir, name='_top_'+str(self.k_final), datarootA=datarootA, datarootB=datarootB)
# self.save_correspondence_as_txt(top_k_correspondence, name='_top_'+str(self.k_final), datarootA=datarootA, datarootB=datarootB)
# with open(os.path.join(self.save_dir, 'correspondence_activation_top' + '.txt'), 'wt') as opt_file:
# for i in range(len(top_k_correspondence[2])):
# opt_file.write('%f\n' % (top_k_correspondence[2][i]))
return scaled_correspondence
def run(self, A, B):
assert (A.size() == B.size())
image_width = A.size(3)
print("Saving original images...")
util.save_final_image(A, 'original_A', self.save_dir)
util.save_final_image(B, 'original_B', self.save_dir)
self.A = self.Tensor(A.size()).copy_(A)
self.B = self.Tensor(B.size()).copy_(B)
print("Starting algorithm...")
# coarse-to-fine
L_start = 5
# level5: 获得 A/B 特征图
self.model.set_input(self.A)
F_A = self.model.forward(level=L_start).data # 相当于 detach 出变量，没有 grad_fn
self.model.set_input(self.B)
F_B = self.model.forward(level=L_start).data
F_Am = F_A.clone()
F_Bm = F_B.clone()
# 1.初始坐标位 a_to_b, b_to_a
initial_map_a_to_b = self.identity_map(F_B.size()) # 1,512,14,14
initial_map_b_to_a = initial_map_a_to_b.clone()
# L: deepest to shallowest
for L in range(L_start, self.L_final - 1, -1): # 5, self.L_final = 2 if fast else 1
patch_size = self.patch_size_list[L - 1] # [3,3]
search_box_radius = self.search_box_radius_list[L - 1] # 2, (path+1)/2
draw_radius = self.draw_radius[L - 1] # 8
if L == L_start:
deepest_level = True # 最深层
correspondence = [] # begin
else:
deepest_level = False
# correspondence: 中间层，基于上轮(+1层) correspondence 寻找 search region
print("Finding best-buddies for the " + str(L) + "-th level")
# 2.对应后坐标位 和 NBBs
# correspondence: [[A_pts], [B_pts], [rank_activation]]
# mapping_a_to_b, mapping_b_to_a; 坐标位
correspondence, mapping_a_to_b, mapping_b_to_a = self.find_neural_best_buddies(
correspondence, F_A, F_B, F_Am, F_Bm, patch_size,
initial_map_a_to_b, initial_map_b_to_a,
search_box_radius, deepest_level
)
# 3.tau 过滤 H(p)
# correspondence: [[A_pts], [B_pts], [rank_activation >= tau]]
correspondence = self.threshold_response_correspondence(correspondence,
self.H_A, self.H_B, # l2 and min_max norm feature
self.tau) # 可以不用手工设置，保留一些统计结果?
print(f'thre >= {self.tau}:', len(correspondence[0]))
# 设置 k_per_level, 由深到浅的每一层 都执行 k-means 选择 NBBs; 减少计算量
if self.k_per_level < float('inf'):
correspondence = self.top_k_in_clusters(correspondence, int(self.k_per_level))
print(f'cluster k({self.k_per_level}):', len(correspondence[0]))
# scale_correspondence 坐标位，在原图画出对应
if L > self.L_final:
print("Drawing correspondence...")
scaled_correspondence = self.scale_correspondence(correspondence, L)
draw.draw_correspondence(self.A, self.B, scaled_correspondence, draw_radius, self.save_dir, L)
F_A, F_B, F_Am, F_Bm, initial_map_a_to_b, initial_map_b_to_a = self.transfer_style_local(
F_A, F_B, patch_size, image_width, mapping_a_to_b, mapping_b_to_a, L
)
filtered_correspondence = self.finalize_correspondence(correspondence, image_width, self.L_final)
draw.draw_correspondence(self.A, self.B, filtered_correspondence, self.draw_radius[self.L_final - 1], self.save_dir)
self.save_correspondence_as_txt(filtered_correspondence)
# After L=1, chose top K by k_final
top_k_correspondence = self.top_k_in_clusters(filtered_correspondence, self.k_final)
draw.draw_correspondence(self.A, self.B, top_k_correspondence, self.draw_radius[self.L_final - 1], self.save_dir, self.k_final)
self.save_correspondence_as_txt(top_k_correspondence, name='_top_' + str(self.k_final))
return scaled_correspondence