def main():
print("started")
config = ConfigProvider.config()
# read data
inspected = cv2.imread(config.data.defective_inspected_path1, 0).astype('float32')
reference = cv2.imread(config.data.defective_reference_path1, 0).astype('float32')
# registration
aligner = Aligner()
warped, warp_mask = aligner.align_images(static=inspected, moving=reference)
# find defects
defect_segmenter = DefectSegmenter()
defect_mask = defect_segmenter.segment_defects(inspected, warped, warp_mask)
# observe results
diff = np.zeros(inspected.shape, dtype=np.float32)
diff[warp_mask] = (np.abs((np.float32(warped) - np.float32(inspected))))[warp_mask]
noise_cleaner = NoiseCleaner()
diff = noise_cleaner.clean_frame(diff, warp_mask)
cv2.imshow("color_result", get_color_diff_image(inspected, defect_mask * 255).astype('uint8'))
plt.imshow(diff.astype('uint8'), cmap='gray')
plt.title("diff")
cv2.imshow("inspected", inspected.astype('uint8'))
cv2.imshow("reference", reference.astype('uint8'))
cv2.imshow("result", defect_mask.astype('uint8') * 255)
plt.show()
cv2.waitKey(0)
class BlurredDiffSegmenter(BaseSegmenter):
"""
Use lower threshold, but lose some accuracy due to bluring.
"""
def __init__(self):
self._config = ConfigProvider.config()
self._noise_cleaner = NoiseCleaner()
self._blured_diff_thres = self._config.detection.blured_diff_thres
@overrides
def detect(self, inspected, warped, warp_mask):
diff = np.zeros(inspected.shape, dtype=np.float32)
diff[warp_mask] = (np.abs(
(np.float32(warped) - np.float32(inspected))))[warp_mask]
diff = self._noise_cleaner.clean_frame(diff, warp_mask)
# make sure noise is not interfering
diff_blured = self._noise_cleaner.blur(diff, sigma=7)
detection_mask = self._blured_diff_thres < diff_blured
# plots
plot_image(diff, "diff")
show_color_diff(warped, inspected, "color diff")
plot_image(diff_blured, "diff_blured")
plot_image(detection_mask, "diff_based_segmentation")
return detection_mask
def __init__(self):
self._config = ConfigProvider.config()
self._noise_cleaner = NoiseCleaner()
self._min_diff_threshold = self._config.refinement.min_diff_threshold
self._dilation_diameter = self._config.refinement.dilation_diameter
self._min_new_connected_component_size = self._config.refinement.min_new_connected_component_size
class DefectSegmentationRefiner(object):
def __init__(self):
self._config = ConfigProvider.config()
self._noise_cleaner = NoiseCleaner()
self._min_diff_threshold = self._config.refinement.min_diff_threshold
self._dilation_diameter = self._config.refinement.dilation_diameter
self._min_new_connected_component_size = self._config.refinement.min_new_connected_component_size
def refine_segmentation(self, focused_defect_mask, inspected, warped,
warp_mask):
"""
see also Segmenter.infer_region_statistics() for more details.
Following commented is the start of a solution that I didn't have time to complete, but is another idea:
After a good segmentation, statistics of pixel color per segment type (there are 3) can be obtained.
Then, each pixel can be classified as defect/good by its color relative to its segment.
Early research showed there are 3 gaussian peaks, roughly centered at gray levels 50, 100, 150.
I didn't have time to create a full working solution like this because my segmentation was lacking.
max_value_per_pixel = np.zeros_like(diff)
min_value_per_pixel = np.zeros_like(diff)
for c in range(config.segmentation.num_classes):
max_value_per_pixel[warped_segmented == c] = statistics_per_class[c][0] + 2 * statistics_per_class[c][1]
min_value_per_pixel[warped_segmented == c] = statistics_per_class[c][0] - 2 * statistics_per_class[c][1]
clean_defect_mask = np.zeros_like(dirty_defect_mask)
clean_defect_mask[dirty_defect_mask] = max_value_per_pixel[dirty_defect_mask] < inspected[dirty_defect_mask]
clean_defect_mask[dirty_defect_mask] = np.logical_or(clean_defect_mask[dirty_defect_mask], min_value_per_pixel[dirty_defect_mask] < inspected[dirty_defect_mask])
plot_image(max_value_per_pixel, "max_value_per_pixel")
plot_image(min_value_per_pixel, "min_value_per_pixel")
"""
diff = np.zeros(inspected.shape, dtype=np.float32)
diff[warp_mask] = (np.abs(
(np.float32(warped) - np.float32(inspected))))[warp_mask]
diff = self._noise_cleaner.clean_frame(diff, warp_mask)
# enlarge detection area in case of close proximity misses
dilated_defect_mask = \
self._noise_cleaner.dilate(focused_defect_mask.copy().astype('uint8'), diameter=self._dilation_diameter).astype(np.bool)
diff_above_thres_mask = self._min_diff_threshold < diff
dilated_defect_mask_with_artifacts = np.logical_and(
diff_above_thres_mask, dilated_defect_mask)
# clean stray pixels which were added due to the dilation, and passed the threshold
enlarged_defect_mask_too_clean = self._noise_cleaner.clean_stray_pixels_bw(
dilated_defect_mask_with_artifacts,
min_size=self._min_new_connected_component_size)
# but not ones that were present before dilation
clean_defect_mask = np.logical_or(enlarged_defect_mask_too_clean,
focused_defect_mask)
plot_image(inspected, "inspected")
plot_image(focused_defect_mask, "focused_defect_mask")
plot_image(clean_defect_mask, "clean_defect_mask")
return clean_defect_mask
def __init__(self):
self._config = ConfigProvider.config()
self._noise_cleaner = NoiseCleaner()
self._thread_defect_high_pass_thres = self._config.detection.thread_defect_high_pass_thres
self._aura_radius = self._config.detection.aura_radius
self._low_diff_far_from_edge_thres = self._config.detection.low_diff_far_from_edge_thres
self._min_thread_defect_size = self._config.detection.min_thread_defect_size
def __init__(self):
self._config = ConfigProvider.config()
self._is_force_translation = self._config.alignment.is_force_translation
self._subpixel_accuracy_resolution = self._config.alignment.subpixel_accuracy_resolution
self._detector = cv2.ORB_create()
self._matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
self._noise_cleaner = NoiseCleaner()
def __init__(self):
self._config = ConfigProvider.config()
self._num_classes = self._config.segmentation.num_classes
self._low_threshold = self._config.segmentation.low_threshold
self._high_threshold = self._config.segmentation.high_threshold
self._auto_thresholds = self._config.segmentation.auto_thresholds
self._noise_cleaner = NoiseCleaner()
self._kmeans = KMeans(n_clusters=self._num_classes)
def _observe_results(defect_mask, inspected, reference, warp_mask, warped):
diff = np.zeros(inspected.shape, dtype=np.float32)
diff[warp_mask] = (np.abs((np.float32(warped) - np.float32(inspected))))[warp_mask]
noise_cleaner = NoiseCleaner()
diff = noise_cleaner.clean_frame(diff, warp_mask)
# cv2.imshow("color_result", get_color_diff_image(inspected, defect_mask * 255).astype('uint8'))
# plt.imshow(diff.astype('uint8'), cmap='gray')
# plt.title("diff")
# cv2.imshow("inspected", inspected.astype('uint8'))
# cv2.imshow("reference", reference.astype('uint8'))
# cv2.imshow("result", defect_mask.astype('uint8') * 255)
plt.show()
cv2.waitKey(0)
class DiffSegmenter(BaseSegmenter):
"""
simplest segmenter, by using a high threshold for noise
"""
def __init__(self):
self._config = ConfigProvider.config()
self._noise_cleaner = NoiseCleaner()
self._diff_thres = self._config.detection.diff_thres
@overrides
def detect(self, inspected, warped, warp_mask):
diff = np.zeros(inspected.shape, dtype=np.float32)
diff[warp_mask] = (np.abs(
(np.float32(warped) - np.float32(inspected))))[warp_mask]
diff = self._noise_cleaner.clean_frame(diff, warp_mask)
detection_mask = self._diff_thres < diff
# plots
plot_image(diff, "diff")
show_color_diff(warped, inspected, "color diff")
plot_image(detection_mask, "diff_based_segmentation")
return detection_mask
class ThreadDefectSegmenter(BaseSegmenter):
"""
find "thin" defects, which are not very different from the background, and not in existing edges proximity.
"""
def __init__(self):
self._config = ConfigProvider.config()
self._noise_cleaner = NoiseCleaner()
self._thread_defect_high_pass_thres = self._config.detection.thread_defect_high_pass_thres
self._aura_radius = self._config.detection.aura_radius
self._low_diff_far_from_edge_thres = self._config.detection.low_diff_far_from_edge_thres
self._min_thread_defect_size = self._config.detection.min_thread_defect_size
@overrides
def detect(self, inspected, warped, warp_mask):
i_blured = self._noise_cleaner.blur(inspected, sigma=10)
high_pass = np.abs(np.float32(i_blured) - np.float32(inspected))
plot_image(i_blured, "i_blured")
plot_image(inspected, "inspected")
show_color_diff(i_blured, inspected, "high_pass")
# find and dilate edges, to get rid of high diff in high pass image caused by real edges
# this will leave us only with edges caused by stuff that weren't in the original image.
# obvious limitation: won't find items near real edges in original image.
edges = cv2.Canny(warped.astype('uint8'), 100, 200) > 0
edges_dialated = self._noise_cleaner.dilate(edges.astype(np.float32),
self._aura_radius)
# blur to avoid noise
high_pass_no_real_edges = high_pass.copy()
high_pass_no_real_edges[~warp_mask] = 0
high_pass_no_real_edges[edges_dialated > 0] = 0
plot_image(high_pass_no_real_edges, "high_pass_no_real_edges")
thread_defect_mask_noisy = self._thread_defect_high_pass_thres < high_pass_no_real_edges
# here we have some false positives, which are caused by noise.
# this detector finds "thread-like" defects, so I require the defects to be connected and at some min size.
thread_defect_mask_clean = self._noise_cleaner.clean_stray_pixels_bw(
thread_defect_mask_noisy, min_size=self._min_thread_defect_size)
thread_defect_mask_closure = self._noise_cleaner.close(
thread_defect_mask_clean.astype('uint8'), diameter=3, iterations=1)
plot_image(thread_defect_mask_noisy, "thread_defect_mask_noisy")
plot_image(thread_defect_mask_clean, "thread_defect_mask_clean")
plot_image(thread_defect_mask_closure, "thread_defect_mask_closure")
return thread_defect_mask_clean
class LowDiffFarFromEdgeSegmenter(BaseSegmenter):
"""
find defects using a very low diff threshold,
but only in case they are far enough from edges in the reference (warped)
"""
def __init__(self):
self._config = ConfigProvider.config()
self._noise_cleaner = NoiseCleaner()
self._aura_radius = self._config.detection.aura_radius
self._low_diff_far_from_edge_thres = self._config.detection.low_diff_far_from_edge_thres
@overrides
def detect(self, inspected, warped, warp_mask):
diff = np.zeros(inspected.shape, dtype=np.float32)
diff[warp_mask] = (np.abs(
(np.float32(warped) - np.float32(inspected))))[warp_mask]
diff = self._noise_cleaner.clean_frame(diff, warp_mask)
# find and dilate edges, to get rid of high diff caused by suboptimal registration
# and is most apparent near real edges.
edges = cv2.Canny(warped.astype('uint8'), 100, 200) > 0
edges_dialated = self._noise_cleaner.dilate(edges.astype(np.float32),
self._aura_radius)
# blur to avoid noise
diff_blured = self._noise_cleaner.blur(diff.copy(), sigma=5)
#disregard area of edges
diff_no_edges_blured = diff_blured
diff_no_edges_blured[edges_dialated > 0] = 0
low_diff_far_from_edge_mask = self._low_diff_far_from_edge_thres < diff_no_edges_blured
plot_image(edges, "edges")
plot_image(edges_dialated, "edges_dilated")
plot_image(diff_no_edges_blured, "diff_no_edges_blured")
plot_image(low_diff_far_from_edge_mask, "low_diff_far_from_edge_mask")
return low_diff_far_from_edge_mask
class Segmenter(object):
"""
This class was another attempt at the problem, which I didn't have time to complete.
Main entry point is infer_region_statistics(), see doc there.
"""
def __init__(self):
self._config = ConfigProvider.config()
self._num_classes = self._config.segmentation.num_classes
self._low_threshold = self._config.segmentation.low_threshold
self._high_threshold = self._config.segmentation.high_threshold
self._auto_thresholds = self._config.segmentation.auto_thresholds
self._noise_cleaner = NoiseCleaner()
self._kmeans = KMeans(n_clusters=self._num_classes)
def _perform_thresholding(self, image):
clean = image.copy()
under_low_mask = clean < self._low_threshold
above_high_mask = self._high_threshold < clean
middle_mask = ~np.logical_or(under_low_mask, above_high_mask)
clean[under_low_mask] = 0
clean[middle_mask] = 1
clean[above_high_mask] = 2
return clean
@staticmethod
def _find_local_minima(a):
mins_mask = np.r_[True, a[1:] < a[:-1]] & np.r_[a[:-1] < a[1:], True]
mins_inds = np.where(mins_mask)[0]
mins = a[mins_inds]
return mins, mins_inds
@staticmethod
def _smooth_curve(curve, kernel_size=5):
kernel = np.ones((kernel_size, )) / kernel_size
smooth = np.convolve(curve, kernel, mode='same')
return smooth
def _infer_thresholds_by_histogram(self, image):
"""
assuming 3 intensity levels, and a valid input image (non defective)
:return: low and high thresholds by which the image can be roughly segmented.
"""
n_bins = 256
threshold_factor = (256 // n_bins)
hist, bins = np.histogram(image.flatten(), bins=n_bins)
smooth_hist = hist
smooth_hist = self._smooth_curve(smooth_hist, 20 // threshold_factor)
smooth_hist = self._smooth_curve(smooth_hist, 14 // threshold_factor)
smooth_hist = self._smooth_curve(smooth_hist, 10 // threshold_factor)
smooth_hist = self._smooth_curve(smooth_hist, 6 // threshold_factor)
# plt.figure()
# plt.plot(hist, color="blue")
# plt.plot(smooth_hist, color="red")
# plt.show()
mins, mins_inds = self._find_local_minima(smooth_hist)
sorted_inds = sorted(
filter(
lambda x: 20 / threshold_factor < x < n_bins - 20 /
threshold_factor, mins_inds))
# assert len(sorted_inds) == self._num_classes - 1
return sorted_inds[0] * threshold_factor, sorted_inds[
1] * threshold_factor, hist, smooth_hist
def segment_image_by_kmeans(self, image):
clean = image
clean = self._noise_cleaner.equalize_histogram(clean)
clean = self._noise_cleaner.clean_salt_and_pepper(clean)
segmentation_map = np.reshape(
self._kmeans.fit_predict(clean.reshape(-1, 1)),
clean.shape).astype(np.uint8)
return segmentation_map
def segment_image_by_threshold(self, image):
clean = image.copy()
clean = self._noise_cleaner.equalize_histogram(clean)
clean = self._noise_cleaner.clean_salt_and_pepper(clean)
hist, smooth_hist = None, None
if self._auto_thresholds:
low_thres, high_thres, hist, smooth_hist = self._infer_thresholds_by_histogram(
clean)
# I allow the high thres to go down, and the low to go up, but not vice versa.
# self._low_threshold, self._high_threshold = max(low_thres, self._low_threshold), min(high_thres, self._high_threshold)
self._low_threshold, self._high_threshold = low_thres, high_thres
segmentation_map = self._perform_thresholding(clean)
return segmentation_map, hist, smooth_hist, self._low_threshold, self._high_threshold
def infer_region_statistics(self, image, mask):
"""
This was supposed to be used for determining 3 regions in the warped image, then calculate the color
distribution within each segment, then know for each pixel its probability of being an outlier by color.
Unfortunately, I didn't have the time to create a good enough segmentation per class and this was no good.
"""
segment_image = self.segment_image_by_kmeans(image.astype('uint8'))
segment_image[~mask] = self._config.segmentation.num_classes
statistics_per_class = []
for c in range(self._config.segmentation.num_classes):
class_data = image[segment_image == c]
m, s = class_data.mean(), class_data.std()
statistics_per_class.append((m, s))
plot_image(image, "image")
plot_image(segment_image, "segment_image")
return statistics_per_class, segment_image
class Aligner(object):
def __init__(self):
self._config = ConfigProvider.config()
self._is_force_translation = self._config.alignment.is_force_translation
self._subpixel_accuracy_resolution = self._config.alignment.subpixel_accuracy_resolution
self._detector = cv2.ORB_create()
self._matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
self._noise_cleaner = NoiseCleaner()
def align_using_feature_matching(self, static, moving):
# this works, but since we know the shift is translation only, normxcorr is better
kp1, des1 = self._detector.detectAndCompute(static, None)
kp2, des2 = self._detector.detectAndCompute(moving, None)
matches = self._matcher.match(des1, des2)
assert len(matches) >= 4 # for perspective tform
moving_pts = np.float32([kp1[m.queryIdx].pt
for m in matches]).reshape(-1, 1, 2)
static_pts = np.float32([kp2[m.trainIdx].pt
for m in matches]).reshape(-1, 1, 2)
tform, mask = cv2.findHomography(moving_pts, static_pts, cv2.RANSAC,
5.0)
tform = self._force_translation_only(tform)
matches_mask = mask.ravel().tolist()
matches_image = cv2.drawMatches(
static,
kp1,
moving,
kp2,
matches,
None,
flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS,
matchesMask=matches_mask,
matchColor=(0, 255, 0),
)
itform = np.linalg.pinv(tform)
return matches_image, itform
def align_using_tform(self, static, moving, tform):
warped = cv2.warpPerspective(moving, tform,
(static.shape[1], static.shape[0]))
warped_region_mask = self._find_warped_region(tform, moving, static)
return warped, warped_region_mask
def align_using_shift(self, static, moving, shift_xy):
tform = np.eye(3)
tform[0, 2] = shift_xy[1]
tform[1, 2] = shift_xy[0]
warped, warped_region_mask = self.align_using_tform(
static, moving, tform)
return warped, warped_region_mask
@staticmethod
def _find_warped_region(tform, moving, static):
white = np.ones(moving.shape)
warped = cv2.warpPerspective(white, tform,
(static.shape[1], static.shape[0]))
warped_region_mask = warped > 0
return warped_region_mask
def _force_translation_only(self, tform):
if self._is_force_translation:
tform[0, 1] = 0
tform[1, 0] = 0
tform[2, 0] = 0
tform[2, 1] = 0
tform[0, 0] = 1
tform[1, 1] = 1
print(f"forcing tform {tform} to translation only")
return tform
@staticmethod
def normxcorr2(template, image, mode="full"):
"""
credit: https://github.com/Sabrewarrior/normxcorr2-python/blob/master/normxcorr2.py
Input arrays should be floating point numbers.
:param template: N-D array, of template or filter you are using for cross-correlation.
Must be less or equal dimensions to image.
Length of each dimension must be less than length of image.
:param image: N-D array
:param mode: Options, "full", "valid", "same"
full (Default): The output of fftconvolve is the full discrete linear convolution of the inputs.
Output size will be image size + 1/2 template size in each dimension.
valid: The output consists only of those elements that do not rely on the zero-padding.
same: The output is the same size as image, centered with respect to the ‘full’ output.
:return: N-D array of same dimensions as image. Size depends on mode parameter.
"""
# If this happens, it is probably a mistake
if np.ndim(template) > np.ndim(image) or \
len([i for i in range(np.ndim(template)) if template.shape[i] > image.shape[i]]) > 0:
print(
"normxcorr2: TEMPLATE larger than IMG. Arguments may be swapped."
)
template = template - np.mean(template)
image = image - np.mean(image)
a1 = np.ones(template.shape)
# Faster to flip up down and left right then use fftconvolve instead of scipy's correlate
ar = np.flipud(np.fliplr(template))
out = fftconvolve(image, ar.conj(), mode=mode)
image = fftconvolve(np.square(image), a1, mode=mode) - \
np.square(fftconvolve(image, a1, mode=mode)) / (np.prod(template.shape))
# Remove small machine precision errors after subtraction
image[np.where(image < 0)] = 0
template = np.sum(np.square(template))
out = out / np.sqrt(image * template)
# Remove any divisions by 0 or very close to 0
out[np.where(np.logical_not(np.isfinite(out)))] = 0
return out
def align_using_normxcorr(self, static, moving):
"""
normxcorr is same as
# conv_res = convolve2d(static, moving, mode='full')
# best_location_conv = np.unravel_index(np.argmax(res), res.shape)
but fast.
"""
res = self.normxcorr2(static, moving, mode='full')
best_location = np.unravel_index(np.argmax(res), res.shape)
moving_should_be_strided_by = -(np.array(best_location) -
np.array(moving.shape) + 1)
return moving_should_be_strided_by
@staticmethod
def align_using_ecc(static, moving):
number_of_iterations = 100
termination_eps = 1e-10
cc, tform = cv2.findTransformECC(
static,
moving,
np.eye(3, dtype=np.float32)[:2, :],
cv2.MOTION_TRANSLATION,
criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT,
number_of_iterations, termination_eps),
inputMask=np.ones(static.shape, dtype='uint8') * 255,
gaussFiltSize=11)
return tform
def align_images(self, static, moving):
# clean noise
static_clean = self._noise_cleaner.clean_salt_and_pepper(static, 5)
moving_clean = self._noise_cleaner.clean_salt_and_pepper(moving, 5)
# enlarge to obtain subpixel accuracy
static_enlarged = cv2.resize(static_clean, (0, 0),
fx=self._subpixel_accuracy_resolution,
fy=self._subpixel_accuracy_resolution)
moving_enlarged = cv2.resize(moving_clean, (0, 0),
fx=self._subpixel_accuracy_resolution,
fy=self._subpixel_accuracy_resolution)
# normxcorr alignment (translation only)
moving_should_be_strided_by_factored = self.align_using_normxcorr(
static=static_enlarged, moving=moving_enlarged)
# return to normal size of translation
moving_should_be_strided_by = np.array(
moving_should_be_strided_by_factored
) / self._subpixel_accuracy_resolution
# perform actual warp
warped, warp_mask = self.align_using_shift(
static.copy(), moving.copy(), moving_should_be_strided_by)
# show result
diff = np.zeros(static.shape, dtype=np.float32)
diff[warp_mask] = (np.abs(
(np.float32(warped) - np.float32(static))))[warp_mask]
show_color_diff(warped, static, "registration")
plot_image(diff, "diff")
return warped, warp_mask
def __init__(self):
self._config = ConfigProvider.config()
self._noise_cleaner = NoiseCleaner()
self._aura_radius = self._config.detection.aura_radius
self._low_diff_far_from_edge_thres = self._config.detection.low_diff_far_from_edge_thres
def __init__(self):
self._config = ConfigProvider.config()
self._noise_cleaner = NoiseCleaner()
self._blured_diff_thres = self._config.detection.blured_diff_thres