-
Notifications
You must be signed in to change notification settings - Fork 12
/
gen_diff_patch.py
114 lines (92 loc) · 4.63 KB
/
gen_diff_patch.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
import os
from os.path import join
import argparse
import sys
import numpy as np
from tqdm import tqdm
sys.path.append('.')
import imageio
import cv2
import matplotlib.pyplot as plt
import json
from matplotlib import cm
def is_overlapping(box1, box2, patch_size):
x1, y1 = box1
x2, y2 = box2
overlap_x = (x1 < (x2 + patch_size)) and (x2 < (x1 + patch_size))
overlap_y = (y1 < (y2 + patch_size)) and (y2 < (y1 + patch_size))
return overlap_x and overlap_y
def find_most_dissimilar_patches(img1, img2, patch_size=128, num_patches=4):
if img1.shape != img2.shape:
raise ValueError("The two images must have the same shape.")
diff = np.abs(img1 - img2)
height, width = img1.shape[:2]
best_patches = []
best_scores = []
candidate_patches = []
for i in tqdm(range(0, height - patch_size), desc='traverse the image grid'):
for j in range(0, width - patch_size):
patch = diff[i:i+patch_size, j:j+patch_size]
score = np.mean(patch)
candidate_patches.append(((i, j), score))
candidate_patches.sort(key=lambda x: x[1], reverse=True)
for (i, j), score in tqdm(candidate_patches, desc='traverse patches'):
if len(best_patches) >= num_patches:
break
if all([not is_overlapping((i, j), best_i_j, patch_size) for best_i_j, _ in best_patches]):
best_patches.append(((i, j), score))
return best_patches
# # 读取图像
# img1 = cv2.imread('image1.jpg', cv2.IMREAD_GRAYSCALE)
# img2 = cv2.imread('image2.jpg', cv2.IMREAD_GRAYSCALE)
# best_patches, best_scores = find_most_dissimilar_patches(img1, img2, patch_size=128, num_patches=4)
# print("Most dissimilar patches at:")
# for (i, j), score in zip(best_patches, best_scores):
# print(f"Top-left corner: ({i}, {j}), Score: {score}")
# cv2.rectangle(img1, (j, i), (j+128, i+128), (255, 0, 0), 2)
# cv2.rectangle(img2, (j, i), (j+128, i+128), (255, 0, 0), 2)
# cv2.imshow('Image 1', img1)
# cv2.imshow('Image 2', img2)
# cv2.waitKey(0)
# cv2.destroyAllWindows()
def vis_patches(img1, img2, patches, patch_size=128, save_path=None):
prop_cycle = plt.rcParams['axes.prop_cycle']
colors = prop_cycle.by_key()['color']
colors_float = [(r, g, b) for r, g, b, _ in [cm.colors.to_rgba(c) for c in colors]]
for idx, ((i, j), score) in enumerate(patches):
color = colors_float[idx % len(colors_float)]
cv2.rectangle(img1, (j, i), (j+patch_size, i+patch_size), color, 2)
cv2.rectangle(img2, (j, i), (j+patch_size, i+patch_size), color, 2)
imageio.imwrite(save_path, np.concatenate([img1, img2], axis=1))
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument('--input', type=str, default='/mnt/data2/home/linhaotong/Datasets/dynerf/cut_beef')
parser.add_argument('--output', type=str, default='/mnt/data2/home/linhaotong/Datasets/dynerf/cut_beef/im4d')
parser.add_argument('--test_view', type=str, default='00')
parser.add_argument('--frame_len', type=int, default=60)
parser.add_argument('--eval_interval', type=int, default=30)
parser.add_argument('--eval_ratio', type=float, default=0.5)
parser.add_argument('--patch_size', type=int, default=128)
parser.add_argument('--num_patches', type=int, default=4)
args = parser.parse_args()
return args
def main(args):
imgs = []
for i in tqdm(range(args.frame_len)):
img_path = join(args.input, 'images', args.test_view, '{:04d}.jpg'.format(i))
img = (imageio.imread(img_path) / 255.).astype(np.float32)
imgs.append(img)
os.makedirs(join(args.output, 'test', 'patches'), exist_ok=True)
os.makedirs(join(args.output, 'test', 'vis'), exist_ok=True)
for i in tqdm(range(0, args.frame_len, args.eval_interval)):
img = cv2.resize(imgs[i], None, fx=args.eval_ratio, fy=args.eval_ratio, interpolation=cv2.INTER_AREA)
avg_img = cv2.resize(np.mean(imgs[i:i+args.eval_interval], axis=0), None, fx=args.eval_ratio, fy=args.eval_ratio, interpolation=cv2.INTER_AREA)
patches = find_most_dissimilar_patches(img, avg_img, patch_size=args.patch_size, num_patches=args.num_patches)
vis_patches(img, avg_img, patches, patch_size=args.patch_size, save_path=join(args.output, 'test', 'vis', '{:03d}.png'.format(i)))
save_dict = {'coords': [patch[0] for patch in patches], 'patch_size': args.patch_size}
save_path = join(args.output, 'test', 'patches', '{:03d}.json'.format(i))
with open(save_path, 'w', encoding='utf-8') as json_file:
json.dump(save_dict, json_file, ensure_ascii=False, indent=4)
if __name__ == '__main__':
args = parse_args()
main(args)