Skip to content
Snippets Groups Projects
Commit 2c98aee0 authored by Alexander Gehrke's avatar Alexander Gehrke
Browse files

wip

parent c19905bc
No related branches found
Tags 1.0.0
No related merge requests found
Hide whitespace changes
Inline Side-by-side
Showing
with 124 additions and 91 deletions
src/linedetect/main.py
+ 124
− 91
View file @ 2c98aee0
from __future__ import annotations
from dataclasses import dataclass
from math import atan2, pi
from sys import argv
from ocr4all.files import imread_bin
from typing import List, TypeVar, Callable, Tuple
import cv2
import PIL
import numpy as np
from dataclasses import dataclass
from typing import List, TypeVar, Callable, Tuple
from functools import partial
from PIL import Image, ImageDraw
PALETTE = (
(1, 0, 0),
......@@ -36,7 +36,26 @@ PALETTE = (
def main():
for file in argv[1:]:
detect_lines(file)
blurred, bin, image = read_and_clean_image(file)
hor_groups, vert_groups = detect_lines(blurred)
conlines = connected_lines(hor_groups, vert_groups, bin.shape[:2])
boxes = Image.new('RGBA', image.size, (255, 0, 0, 0))
draw = ImageDraw.Draw(boxes)
for component in conlines.components():
color = (255, 0, 0, 128) if is_boxed(component) else (0, 255, 0, 128)
draw.rectangle(component.to_rect(), fill=color)
lines_view = line_image(hor_groups, vert_groups, bin.shape[:2])
out = Image.alpha_composite(image.convert('RGBA'), boxes)
out.save(f'{file}.out.png')
#show_side_by_side(out, lines_view)
def is_boxed(component: ConnectedComponent) -> bool:
cutout = component.cutout()
cutout[10:-10, 10:-10] = 0
return np.sum(cutout) > 8 * np.sum(cutout.shape)
@dataclass
......@@ -45,6 +64,17 @@ class ConnectedComponent:
area: int
width: int
height: int
pos: Tuple[int, int]
parent: ConnectedComponents
def to_rect(self):
return (self.pos[0], self.pos[1], self.pos[0] + self.width, self.pos[1] + self.height)
def cutout(self):
return self.parent.labels[
self.pos[1]:self.pos[1] + self.height,
self.pos[0]:self.pos[0] + self.width
] == self.label
class ConnectedComponents:
......@@ -57,7 +87,12 @@ class ConnectedComponents:
area = self.stats[i, cv2.CC_STAT_AREA]
w = self.stats[i, cv2.CC_STAT_WIDTH]
h = self.stats[i, cv2.CC_STAT_HEIGHT]
yield ConnectedComponent(i, area, w, h)
l = self.stats[i, cv2.CC_STAT_LEFT]
t = self.stats[i, cv2.CC_STAT_TOP]
yield ConnectedComponent(i, area, w, h, (l, t), self)
def __str__(self):
return f'ConnectedComponents(num_labels={self.num_labels}, labels={self.labels}, stats={self.stats}, centroids={self.centroids})'
def find_large_components(image, *, min_area: int = 0, min_length: int = 0, min_length_on_both_axis: bool = False):
......@@ -98,24 +133,86 @@ def group_by_predicate(elements: List[A], pred: Callable[[A, A], bool]) -> List[
return grouped
def detect_lines(file):
def detect_lines(blurred, *, rho=1, theta=np.pi / 180, threshold=15, min_line_length=100, max_line_gap=20):
hor_groups, vert_groups = grouped_hough_lines(blurred, rho=rho, theta=theta, threshold=threshold,
min_line_length=min_line_length, max_line_gap=max_line_gap)
hor_groups = [simplify_group(g, horizontal=True) for g in hor_groups]
vert_groups = [simplify_group(g, horizontal=False) for g in vert_groups]
return hor_groups, vert_groups
def detect_lines_debug(file):
blurred, bin, _ = read_and_clean_image(file)
# rhos = []
# for rho in range(1,5):
# hor_groups, vert_groups = grouped_hough_lines(blurred, rho=rho, theta=np.pi / 180, threshold=15, min_line_length=100,
# max_line_gap=20)
# lines_view = line_image(hor_groups, vert_groups, bin.shape[:2])
# rhos.append(lines_view)
thetas = []
for theta in range(180, 360, 30):
hor_groups, vert_groups = grouped_hough_lines(blurred, rho=1, theta=np.pi / theta, threshold=15,
min_line_length=100,
max_line_gap=20)
lines_view = line_image(hor_groups, vert_groups, bin.shape[:2])
thetas.append(lines_view)
# lengths = []
# for length in range(50,101,10):
# hor_groups, vert_groups = grouped_hough_lines(blurred, rho=1, theta=np.pi / 180, threshold=15, min_line_length=length,
# max_line_gap=20)
# lines_view = line_image(hor_groups, vert_groups, bin.shape[:2])
# lengths.append(lines_view)
# hor_groups = [simplify_group(g, horizontal=True) for g in hor_groups]
# vert_groups = [simplify_group(g, horizontal=False) for g in vert_groups]
# lines_view2 = line_image(hor_groups, vert_groups, bin.shape[:2])
# show_side_by_side(bin, blurred, lines_view, lines_view2)
show_side_by_side(*thetas)
# return hor_groups
def read_and_clean_image(file):
print(f"Detecting lines in {file}")
image = imread_bin(file)
line_min_length = min(image.shape) / 30
(num_labels, cc_labels, large_ccs) = find_large_components(image, min_length=line_min_length)
pil_image = Image.open(file)
# noinspection PyTypeChecker
bin = np.logical_not(np.array(pil_image.convert('1')))
image = pil_image.convert('RGB')
line_min_length = min(bin.shape) / 30
(num_labels, cc_labels, large_ccs) = find_large_components(bin, min_length=line_min_length)
cleaned_image = filter_components(num_labels, cc_labels, large_ccs)
blurred = binary_blur(cleaned_image * 255, kernel_size=15)
print("Cleaned binary image")
return blurred, bin, image
def grouped_hough_lines(bin_image, *, rho=1, theta=np.pi / 180, threshold=15, min_line_length=100, max_line_gap=20):
"""
rho = 1 # distance resolution in pixels of the Hough grid
theta = np.pi / 180 # angular resolution in radians of the Hough grid
threshold = 15 # minimum number of votes (intersections in Hough grid cell)
min_line_length = 100 # minimum number of pixels making up a line
max_line_gap = 20 # maximum gap in pixels between connectable line segments
lines: np.ndarray = cv2.HoughLinesP(blurred, rho, theta, threshold, None,
:param bin_image:
:param rho:
:param theta:
:param threshold:
:param min_line_length:
:param max_line_gap:
:return:
"""
lines: np.ndarray = cv2.HoughLinesP(bin_image, rho, theta, threshold, None,
min_line_length, max_line_gap)
# show_side_by_side(image, blurred)
hor, vert = split_lines_by_direction(lines)
def lines_overlap(a: np.ndarray, b: np.ndarray, vertical: bool, dist_threshold: int = 10):
......@@ -132,15 +229,7 @@ def detect_lines(file):
hor_groups = group_by_predicate(hor, lambda a, b: lines_overlap(a, b, vertical=False, dist_threshold=10))
vert_groups = group_by_predicate(vert, lambda a, b: lines_overlap(a, b, vertical=True, dist_threshold=10))
lines_view = line_image(hor_groups, vert_groups, image.shape[:2])
hor_groups = [simplify_group(g, horizontal=True) for g in hor_groups]
vert_groups = [simplify_group(g, horizontal=False) for g in vert_groups]
lines_view2 = line_image(hor_groups, vert_groups, image.shape[:2])
show_side_by_side(image, blurred, lines_view, lines_view2)
return hor_groups
return hor_groups, vert_groups
def simplify_group(line_group: List[np.ndarray], *, horizontal: bool) -> np.ndarray:
......@@ -175,19 +264,26 @@ def simplify_group(line_group: List[np.ndarray], *, horizontal: bool) -> np.ndar
def line_image(hor, vert, shape):
line_image = np.zeros(shape + (3,))
for index, color_group in enumerate(zip(PALETTE, hor)):
print(f"[{index}] {color_group[0]}")
draw_line_group(line_image, color_group[1], color_group[0], index)
for index, color_group in enumerate(zip(PALETTE, vert)):
print(f"[{index}] {color_group[0]}")
draw_line_group(line_image, color_group[1], color_group[0], index)
return line_image
def connected_lines(hor, vert, shape):
line_image = np.zeros(shape)
for line in hor + vert:
cv2.polylines(line_image, np.array([line]), isClosed=False, color=(1, 1, 1), thickness=5)
return ConnectedComponents(line_image)
def show_side_by_side(*comparison_images: np.ndarray):
print("Showing images")
from matplotlib import pyplot as plt
f, ax = plt.subplots(1, len(comparison_images), sharex='all', sharey='all')
f, ax = plt.subplots(1, len(comparison_images), sharex='all', sharey='all', squeeze=False)
for idx, image in enumerate(comparison_images):
ax[idx].imshow(image)
print(f"Showing image {idx}")
ax[0][idx].imshow(image)
plt.show()
......@@ -233,68 +329,5 @@ def binary_blur(image, *, low_threshold=50, high_threshold=255, kernel_size=5):
return cv2.threshold(blur_gray, low_threshold, high_threshold, cv2.THRESH_BINARY)[1]
def detect_lines_old(file):
print(f"Detecting lines in {file}")
image = imread_bin(file)
output = cv2.connectedComponentsWithStats(image.astype('uint8'), connectivity=8)
(numLabels, labels, stats, centroids) = output
max = 1
max_area = stats[1, cv2.CC_STAT_AREA]
for i in range(2, numLabels):
area = stats[i, cv2.CC_STAT_AREA]
if area > max_area:
max_area = area
max = i
# extract the connected component statistics and centroid for
# the current label
x = stats[max, cv2.CC_STAT_LEFT]
y = stats[max, cv2.CC_STAT_TOP]
w = stats[max, cv2.CC_STAT_WIDTH]
h = stats[max, cv2.CC_STAT_HEIGHT]
area = stats[max, cv2.CC_STAT_AREA]
(cX, cY) = centroids[max]
output = np.array(PIL.Image.open(file).convert('RGB'))
output = output[:, :, ::-1].copy()
cv2.rectangle(output, (x, y), (x + w, y + h), (0, 255, 0), 3)
cv2.circle(output, (int(cX), int(cY)), 4, (0, 0, 255), -1)
componentMask = (labels == max).astype("uint8") * 255
kernel_size = 5
blur_gray = cv2.GaussianBlur(componentMask, (kernel_size, kernel_size), 0)
low_threshold = 127
high_threshold = 255
(_, edges) = cv2.threshold(blur_gray, low_threshold, high_threshold, cv2.THRESH_BINARY)
rho = 1 # distance resolution in pixels of the Hough grid
theta = np.pi / 180 # angular resolution in radians of the Hough grid
threshold = 15 # minimum number of votes (intersections in Hough grid cell)
min_line_length = 100 # minimum number of pixels making up a line
max_line_gap = 20 # maximum gap in pixels between connectable line segments
line_image = np.copy(componentMask) * 0 # creating a blank to draw lines on
# Run Hough on edge detected image
# Output "lines" is an array containing endpoints of detected line segments
lines = cv2.HoughLinesP(edges, rho, theta, threshold, np.array([]),
min_line_length, max_line_gap)
# Draw the lines on the image
lines_edges = cv2.addWeighted(componentMask, 0.8, line_image, 1, 0)
for line in lines:
for x1, y1, x2, y2 in line:
cv2.line(line_image, (x1, y1), (x2, y2), (255, 0, 0), 5)
# show our output image and connected component mask
cv2.imshow("bin", blur_gray)
cv2.imshow("Output", line_image)
# cv2.imshow("Connected Component", componentMask)
while cv2.waitKey(0) not in [27, ord('q')]:
pass
if __name__ == "__main__":
main()
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment