-
Notifications
You must be signed in to change notification settings - Fork 0
/
support.py
183 lines (158 loc) · 6.42 KB
/
support.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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
"""
Support file for the Vesuvius Basic Compression script.
Most of these are wrappers for cv2 functions alongside
some experimental mask generation functions.
Dependencies
------------
Libraries: os, cv2, numpy and collections.
"""
import os
import cv2
import numpy as np
from collections import deque
# Functions to support file handling.
def check_path(path):
if not os.path.isfile(path):
print(f'Error - Source file ({path}) could not be found.')
return False
return True
def load_image(path, grayscale=True):
image = cv2.imread(path, cv2.IMREAD_UNCHANGED)
if image.shape[-1] == 3 and grayscale:
return np.average(image, axis=2)
return image
def save_image(image, path):
cv2.imwrite(path, image)
# Functions to support data structure handling.
def find_value_in_2D(array, value, start=0):
for i in range(start, array.shape[0]):
for j in range(0, array.shape[1]):
if array[i, j] == value:
return [i, j]
# Basic image processing
def clip(image, clip_min, clip_max, skip=False):
if not skip:
image[image<clip_min] = clip_min
image[image>clip_max] = clip_max
return (image-clip_min)/(clip_max-clip_min)
def set_to_int8(image):
return (image*255).astype('uint8')
def blur_image(image, kernel):
return cv2.filter2D(src=image, ddepth=-1, kernel=kernel)
# Functions to support breadth-first search.
def bfs(position, grid):
"""
A Breadth First Search implementation.
Expects a position as [y, x] and a grid to
traverse. The grid should be a 2D numpy array
with values of 1 for any traversable position.
Returns a list of traversed points and an
updated grid (0 is used for traversed points).
"""
queue, closed = deque([position]), []
while queue:
current = queue.popleft()
closed.append(current)
grid[current[0], current[1]] = 0
queue = fetch_neighbours(current, queue, grid)
return closed, grid
def fetch_neighbours(point, queue, grid):
if point[0] > 0 and grid[point[0]-1, point[1]] == 1 and [point[0]-1, point[1]] not in queue:
queue.append([point[0]-1, point[1]])
if point[0] < grid.shape[0]-1 and grid[point[0]+1, point[1]] == 1 and [point[0]+1, point[1]] not in queue:
queue.append([point[0]+1, point[1]])
if point[1] > 0 and grid[point[0], point[1]-1] == 1 and [point[0], point[1]-1] not in queue:
queue.append([point[0], point[1]-1])
if point[1] < grid.shape[1]-1 and grid[point[0], point[1]+1] == 1 and [point[0], point[1]+1] not in queue:
queue.append([point[0], point[1]+1])
return queue;
# Functions to support edge detection.
def sobel_edge_detector(image):
grad_x = cv2.Sobel(image, cv2.CV_64F, 1, 0)
grad_y = cv2.Sobel(image, cv2.CV_64F, 0, 1)
grad = np.sqrt(grad_x**2 + grad_y**2)
grad_norm = (grad * 255 / grad.max()).astype('uint8')
return grad_norm
def clean_edges(image, min_check):
position, grid = [0, 0], image.copy()
grid[grid > 0] = 1
while 1 in grid:
position = find_value_in_2D(array=grid, value=1, start=position[0])
cluster, grid = bfs(position, grid)
if check_cluster(cluster, image.shape[0]-1, image.shape[1]-1):
if min_check and len(cluster) > 1000:
continue
else:
for point in cluster:
image[point[0], point[1]] = 0
return image
def check_cluster(points, max_y, max_x):
for point in points:
if point[0] == 0 or point[0] == max_y or point[1] == 0 or point[1] == max_x:
return False
return True
# Mask creation and handling functions.
def density_blur(image, cycles, small_blur, large_blur, diff_multiplier):
for i in range(0, cycles):
# Initial blur to smooth pixel local pixel differences.
blur_1 = blur_image(image, np.ones((small_blur, small_blur), np.float32)/(small_blur**2))
# Second blur to get averages of larger local environments.
blur_2 = blur_image(image, np.ones((large_blur, large_blur), np.float32)/(large_blur**2))
diff = np.abs(blur_1.astype('int16') - blur_2.astype('int16'))
diff = (image - diff_multiplier*diff)
diff[diff < 0] = 0
image = set_to_int8(clip(diff, diff.min(), diff.max(), True))
return image
def cluster_segmentation(image, num_clusters, grayscale=True):
X = image.reshape(-1, 1).repeat(3, axis=1).astype('float32')
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.2)
_, labels, (centers) = cv2.kmeans(X, num_clusters, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
centers = centers.astype('uint8')
labels = labels.flatten()
segmented = centers[labels.flatten()]
segmented = segmented.reshape((image.shape[0], image.shape[1], 3))
if grayscale:
return segmented[:, :, 0]
return segmented
def set_frame(image, frame):
image[:frame[0]] = 200
image[image.shape[0]-frame[1]:] = 200
image[:, :frame[2]] = 200
image[:, image.shape[1]-frame[3]:] = 200
return image
def create_mask(image, frame, position, min_check):
# Clip image to remove the background noise and increase sheet contrast.
clipped = set_to_int8(clip(image, 30, image.max()))
# Use density difference with blur to highlight the case.
dense = density_blur(clipped, 1, 10, 40, 3)
# Final blur to create more homogenous regions.
blur = blur_image(dense, np.ones((80, 80), np.float32)/6400)
blur = set_to_int8(clip(blur, blur.min(), blur.max(), True))
# Resize for efficient pixel analysis.
small = cv2.resize(blur, (blur.shape[1]//10, blur.shape[0]//10))
# Segmentation by cluster analysis.
mask = cluster_segmentation(small, 3)
# Edge detection.
mask[mask < mask.max()] = 0
mask[mask > 0] = 255
mask = sobel_edge_detector(mask)
mask = clean_edges(mask, min_check)
# Floodfill to create mask.
mask[mask>0] = 200
frame = [i//10 for i in frame]
mask = set_frame(mask, frame)
cv2.floodFill(mask, None, (position[0]//10, position[1]//10), 255)
mask[mask < 255] = 0
# Resize to original image shape.
mask = cv2.resize(mask, (image.shape[1], image.shape[0]))
# Make mask binary.
mask = mask // 255
return mask
def apply_mask(image, mask):
"""
A function to mask out regions of an
image using a binary mask for the image.
This function returns an image array.
"""
image = image * mask
return image