ymcmy commited on
Commit
bad2df4
·
verified ·
1 Parent(s): 0fb03d2

Create utils.py

Browse files
Files changed (1) hide show
  1. utils.py +135 -0
utils.py ADDED
@@ -0,0 +1,135 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import cv2
3
+ import PIL.Image
4
+ from scipy.interpolate import griddata
5
+
6
+ def RGB2gray(rgb):
7
+ r, g, b = rgb[:,:,0], rgb[:,:,1], rgb[:,:,2]
8
+ gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
9
+ return gray
10
+
11
+ # Update img_to_patches to handle direct image input
12
+ def img_to_patches(img: PIL.Image.Image) -> tuple:
13
+ patch_size = 16
14
+ img = img.convert('RGB') # Ensure image is in RGB format
15
+
16
+ grayscale_imgs = []
17
+ imgs = []
18
+ coordinates = []
19
+
20
+ for i in range(0, img.height, patch_size):
21
+ for j in range(0, img.width, patch_size):
22
+ box = (j, i, j + patch_size, i + patch_size)
23
+ img_color = np.asarray(img.crop(box))
24
+ grayscale_image = cv2.cvtColor(src=img_color, code=cv2.COLOR_RGB2GRAY)
25
+ grayscale_imgs.append(grayscale_image.astype(dtype=np.int32))
26
+ imgs.append(img_color)
27
+ normalized_coord = (i + patch_size // 2, j + patch_size // 2)
28
+ coordinates.append(normalized_coord)
29
+
30
+ return grayscale_imgs, imgs, coordinates, (img.height, img.width)
31
+
32
+ def get_l1(v):
33
+ return np.sum(np.abs(v[:, :-1] - v[:, 1:]))
34
+
35
+ def get_l2(v):
36
+ return np.sum(np.abs(v[:-1, :] - v[1:, :]))
37
+
38
+ def get_l3l4(v):
39
+ l3 = np.sum(np.abs(v[:-1, :-1] - v[1:, 1:]))
40
+ l4 = np.sum(np.abs(v[1:, :-1] - v[:-1, 1:]))
41
+ return l3 + l4
42
+
43
+ def get_pixel_var_degree_for_patch(patch: np.array) -> int:
44
+ l1 = get_l1(patch)
45
+ l2 = get_l2(patch)
46
+ l3l4 = get_l3l4(patch)
47
+ return l1 + l2 + l3l4
48
+
49
+ def get_rich_poor_patches(img: PIL.Image.Image, coloured=True):
50
+ gray_scale_patches, color_patches, coordinates, img_size = img_to_patches(img)
51
+ var_with_patch = []
52
+ for i, patch in enumerate(gray_scale_patches):
53
+ if coloured:
54
+ var_with_patch.append((get_pixel_var_degree_for_patch(patch), color_patches[i], coordinates[i]))
55
+ else:
56
+ var_with_patch.append((get_pixel_var_degree_for_patch(patch), patch, coordinates[i]))
57
+
58
+ var_with_patch.sort(reverse=True, key=lambda x: x[0])
59
+ mid_point = len(var_with_patch) // 2
60
+ r_patch = [(patch, coor) for var, patch, coor in var_with_patch[:mid_point]]
61
+ p_patch = [(patch, coor) for var, patch, coor in var_with_patch[mid_point:]]
62
+ p_patch.reverse()
63
+ return r_patch, p_patch, img_size
64
+
65
+ def azimuthalAverage(image, center=None):
66
+ y, x = np.indices(image.shape)
67
+ if not center:
68
+ center = np.array([(x.max() - x.min()) / 2.0, (y.max() - y.min()) / 2.0])
69
+ r = np.hypot(x - center[0], y - center[1])
70
+ ind = np.argsort(r.flat)
71
+ r_sorted = r.flat[ind]
72
+ i_sorted = image.flat[ind]
73
+ r_int = r_sorted.astype(int)
74
+ deltar = r_int[1:] - r_int[:-1]
75
+ rind = np.where(deltar)[0]
76
+ nr = rind[1:] - rind[:-1]
77
+ csim = np.cumsum(i_sorted, dtype=float)
78
+ tbin = csim[rind[1:]] - csim[rind[:-1]]
79
+ radial_prof = tbin / nr
80
+ return radial_prof
81
+
82
+ def azimuthal_integral(img, epsilon=1e-8, N=50):
83
+ if len(img.shape) == 3 and img.shape[2] == 3:
84
+ img = RGB2gray(img)
85
+ f = np.fft.fft2(img)
86
+ fshift = np.fft.fftshift(f)
87
+ fshift += epsilon
88
+ magnitude_spectrum = 20 * np.log(np.abs(fshift))
89
+ psd1D = azimuthalAverage(magnitude_spectrum)
90
+ points = np.linspace(0, N, num=psd1D.size)
91
+ xi = np.linspace(0, N, num=N)
92
+ interpolated = griddata(points, psd1D, xi, method='cubic')
93
+ interpolated = (interpolated - np.min(interpolated)) / (np.max(interpolated) - np.min(interpolated))
94
+ return interpolated.astype(np.float32)
95
+
96
+ def positional_emb(coor, im_size, N):
97
+ img_height, img_width = im_size
98
+ center_y, center_x = coor
99
+ normalized_y = center_y / img_height
100
+ normalized_x = center_x / img_width
101
+ pos_emb = np.zeros(N)
102
+ indices = np.arange(N)
103
+ div_term = 10000 ** (2 * (indices // 2) / N)
104
+ pos_emb[0::2] = np.sin(normalized_y / div_term[0::2]) + np.sin(normalized_x / div_term[0::2])
105
+ pos_emb[1::2] = np.cos(normalized_y / div_term[1::2]) + np.cos(normalized_x / div_term[1::2])
106
+ return pos_emb
107
+
108
+ def azi_diff(img: PIL.Image.Image, patch_num, N):
109
+ r, p, im_size = get_rich_poor_patches(img)
110
+ r_len = len(r)
111
+ p_len = len(p)
112
+ patch_emb_r = np.zeros((patch_num, N))
113
+ patch_emb_p = np.zeros((patch_num, N))
114
+ positional_emb_r = np.zeros((patch_num, N))
115
+ positional_emb_p = np.zeros((patch_num, N))
116
+ coor_r = []
117
+ coor_p = []
118
+ if r_len != 0:
119
+ for idx in range(patch_num):
120
+ tmp_patch1 = r[idx % r_len][0]
121
+ tmp_coor1 = r[idx % r_len][1]
122
+ patch_emb_r[idx] = azimuthal_integral(tmp_patch1, N=N)
123
+ positional_emb_r[idx] = positional_emb(tmp_coor1, im_size, N)
124
+ coor_r.append(tmp_coor1)
125
+ if p_len != 0:
126
+ for idx in range(patch_num):
127
+ tmp_patch2 = p[idx % p_len][0]
128
+ tmp_coor2 = p[idx % p_len][1]
129
+ patch_emb_p[idx] = azimuthal_integral(tmp_patch2, N=N)
130
+ positional_emb_p[idx] = positional_emb(tmp_coor2, im_size, N)
131
+ coor_p.append(tmp_coor2)
132
+ output = {"total_emb": [patch_emb_r + positional_emb_r / 5, patch_emb_p + positional_emb_p / 5],
133
+ "positional_emb": [positional_emb_r / 5, positional_emb_p / 5], "coor": [coor_r, coor_p],
134
+ "image_size": im_size}
135
+ return output