def line_extraction_v1(
probs: np.ndarray,
low_threshold: float,
high_threshold: float,
sigma: float = 0.0,
filter_width: float = 0.00,
vertical_maxima: bool = False) -> Tuple[List[np.ndarray], np.ndarray]:
# Smooth
probs2 = cleaning_probs(probs, sigma=sigma)
lines_mask = hysteresis_thresholding(
probs2,
low_threshold,
high_threshold,
candidates_mask=vertical_local_maxima(probs2)
if vertical_maxima else None)
# Remove lines touching border
# lines_mask = remove_borders(lines_mask)
# Extract polygons from line mask
contours = find_lines(lines_mask)
filtered_contours = []
page_width = probs.shape[1]
for cnt in contours:
centroid_x, centroid_y = np.mean(cnt, axis=0)[0]
if centroid_x < filter_width * page_width or centroid_x > (
1 - filter_width) * page_width:
continue
# if cv2.arcLength(cnt, False) < filter_width*page_width:
# continue
filtered_contours.append(cnt)
return filtered_contours, lines_mask
def line_extraction_v0(probs, sigma, threshold):
# probs_line = probs[:, :, 1]
probs_line = probs
# Smooth
probs2 = cv2.GaussianBlur(probs_line, (int(3*sigma)*2+1, int(3*sigma)*2+1), sigma)
lines_mask = probs2 >= threshold
# Extract polygons from line mask
contours = find_lines(lines_mask)
return contours, lines_mask
def detect(self, img, model_predict_h_w=None):
'''
Input:
img: a BGR image (np.ndarray)
Return:
cv2 style contours
Reference:
https://github.com/dhlab-epfl/fdh-tutorials/blob/master/computer-vision-deep-learning/3-applications/dl-document-processing-textlines/fdh_document_processing.ipynb
'''
assert isinstance(img, np.ndarray) and len(img.shape) == 3
assert model_predict_h_w is None or isinstance(model_predict_h_w,
tuple)
with self.sess.as_default():
# Deep Learning based textline detection
start = time.time()
# Note: the model takes RGB image as input
output_textline = self.model.predict(img[:, :, [2, 1, 0]])
if self.debug:
print("[!] The model took {} to predict this image".format(
time.time() - start))
textline_probs = output_textline['probs'][0, :, :, 1]
if self.debug:
plt.imshow(textline_probs)
plt.savefig(str(self.debug_dir / "textline_probability.png"))
# The higher the sigma, the less number of textlines we have
textline_probs2 = cleaning_probs(textline_probs, sigma=1)
textline_mask = hysteresis_thresholding(textline_probs2,
low_threshold=0.3,
high_threshold=0.6,
candidates_mask=None)
if self.debug:
plt.imshow(textline_mask)
plt.savefig(str(self.debug_dir / "textline_mask.png"))
start = time.time()
line_contours = line_vectorization.find_lines(
resize(textline_mask, img.shape[0:2]))
if self.debug:
print("[!] Find lines took {} secs".format(time.time() -
start))
if self.debug:
drawn_line_img = cv2.drawContours(img.copy(),
line_contours,
-1, (0, 255, 0),
thickness=3)
cv2.imwrite(str(self.debug_dir / "drawn_line_img.png"),
drawn_line_img)
return line_contours
def line_extraction_v1(
probs: np.ndarray,
low_threshold: float,
high_threshold: float,
sigma: float = 0.0,
filter_width: float = 0.00,
vertical_maxima: bool = False) -> Tuple[List[np.ndarray], np.ndarray]:
"""
Given a probability map, returns the contour of lines and the corresponding mask
:param probs: probability map (numpy array)
:param low_threshold: hysteresis low threshold
:param high_threshold: hysteresis high threshold
:param sigma: sigma value for gaussian filtering
:param filter_width: percentage of the image width to filter out lines that are close to borders (default 0.0)
:param vertical_maxima: set to True to use vertical local maxima as candidates for the hysteresis thresholding
:return:
"""
# Smooth
probs2 = cleaning_probs(probs, sigma=sigma)
lines_mask = hysteresis_thresholding(
probs2,
low_threshold,
high_threshold,
candidates_mask=vertical_local_maxima(probs2)
if vertical_maxima else None)
# Remove lines touching border
# lines_mask = remove_borders(lines_mask)
# Extract polygons from line mask
contours = find_lines(lines_mask)
filtered_contours = []
page_width = probs.shape[1]
for cnt in contours:
centroid_x, centroid_y = np.mean(cnt, axis=0)[0]
if centroid_x < filter_width * page_width or centroid_x > (
1 - filter_width) * page_width:
continue
# if cv2.arcLength(cnt, False) < filter_width*page_width:
# continue
filtered_contours.append(cnt)
return filtered_contours, lines_mask