Skip to content

Instantly share code, notes, and snippets.

@petrblahos

petrblahos/hla.py

Created Oct 15, 2018
Embed
What would you like to do?
import sys
import cv2
import numpy as np
def main(fn):
SZ = 21
OUTPUT_ARRAY = [[0]*SZ for i in range(SZ)]
img = cv2.imread(fn)
gr = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
gr = cv2.threshold(gr, 190, 255, cv2.THRESH_BINARY_INV)[1]
kernel = np.ones((2, 2), dtype=np.uint8)
gr = cv2.morphologyEx(gr, cv2.MORPH_OPEN, kernel)
gr = cv2.morphologyEx(gr, cv2.MORPH_CLOSE, kernel)
(__im2, cnts, __hierarchy) = cv2.findContours(gr, cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)
# get the biggest shape
cnts.sort(key=cv2.contourArea)
cnt = cnts[-1]
cv2.drawContours(img, [cnt], -1, (0, 255, 255), 4)
# find the enclosing square
x, y, w, h = cv2.boundingRect(cnt)
cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)
mask = np.zeros(gr.shape, dtype=np.uint8)
# now, we know that the field is 31x31
w0 = w / SZ
h0 = h / SZ
for i in range(SZ):
for j in range(SZ):
cv2.rectangle(
mask,
(int(x + i * w0 + 10), int(y + j * w0 + 12)),
(int(x + i * w0 + w0 - 14), int(y + j * h0 + h0 - 12)),
255, -1,
)
gr = cv2.bitwise_and(gr, mask)
KNOWN_NUMBERS = [
# number, x-coord, y-coord
(1, 9, 2), (1, 10, 16), (1, 14, 12),
(2, 9, 5), (2, 2, 6), (2, 2, 7), (2, 12, 17), (2, 9, 15),
(3, 10, 1), (3, 10, 11), (3, 10, 14), (3, 10, 18), (3, 18, 5),
(3, 2, 4), (3, 16, 9),
(4, 10, 2), (4, 1, 7), (4, 1, 8),
(5, 11, 2), (5, 18, 9), (5, 1, 11), (5, 3, 16),
(6, 13, 1), (6, 9, 16),
(7, 13, 2), (7, 1, 10), (7, 3, 4), (7, 11, 0),
(8, 9, 1), (8, 6, 4), (8, 20, 11), (8, 4, 6),
(9, 6, 3), (9, 6, 12),
]
#serious_mark(gr, KNOWN_NUMBERS, 1, x, y, w0, h0)
TEMPLATES = []
for (n, xx, yy) in KNOWN_NUMBERS:
TEMPLATES.append((n, get_sq_function(gr, xx, yy, x, y, w0, h0)))
# visualize_search(gr, TEMPLATES, x, y, w0, h0, 10, 16)
out = cv2.cvtColor(gr, cv2.COLOR_GRAY2RGB)
for i in range(SZ):
for j in range(SZ):
sq = get_sq_function(gr, i, j, x, y, w0, h0)
if sq is None:
continue
results = []
for (number, templ) in TEMPLATES:
result = cv2.matchTemplate(sq, templ, cv2.TM_CCOEFF)
(_, score, _, _) = cv2.minMaxLoc(result)
if score > 1:
results.append((score, number))
if results:
results.sort()
mark_square(out, i, j, x, y, w0, h0, results[-1][1])
OUTPUT_ARRAY[j][i] = results[-1][1]
cv2.imwrite("out.jpg", out)
return OUTPUT_ARRAY
def serious_mark(img, nums, hl, x, y, w0, h0):
for (n, xx, yy) in nums:
if n != hl:
continue
mark_square(img, xx, yy, x, y, w0, h0, n)
cv2.imwrite("out.jpg", img)
sys.exit(0)
def visualize_search(img, TEMPLATES, x, y, w0, h0, xx, yy):
out = np.zeros((300, 2000), dtype=np.uint8)
sq = get_sq_function(img, xx, yy, x, y, w0, h0)
X = 20
results = []
for (number, templ) in TEMPLATES:
result = cv2.matchTemplate(sq, templ, cv2.TM_CCOEFF)
out[20:20 + sq.shape[0], X:X + sq.shape[1]] = sq
out[150:150 + templ.shape[0], X:X + templ.shape[1]] = templ
(_, score, _, _) = cv2.minMaxLoc(result)
results.append((score, number))
X += sq.shape[1] + 5
best = max([i[0] for i in results])
X = 20
for (val, num) in results:
cv2.putText(out, "%0.2f" % (val / best, ),
(X, 120),
cv2.FONT_HERSHEY_SIMPLEX,
0.4, 255, 1, cv2.LINE_AA)
X += sq.shape[1] + 5
cv2.imwrite("out.jpg", out)
sys.exit(0)
def find_fit_square(img, x, y, x0, y0, w0, h0):
sub_img = extract_square(img, x, y, x0, y0, w0, h0)
(__im2, cnts, __hierarchy) = cv2.findContours(sub_img, cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)
if not cnts:
return None
min_x = min_y = 100000
max_x = max_y = -100000
for i in cnts:
c = i.reshape((i.shape[0] * i.shape[-1]))
min_x = min((min_x, min(c[0::2])))
max_x = max((max_x, max(c[0::2])))
min_y = min((min_y, min(c[1::2])))
max_y = max((max_y, max(c[1::2])))
sub_sub = sub_img[min_y:max_y, min_x:max_y]
res = cv2.resize(sub_sub, (int(w0), int(h0)), interpolation=cv2.INTER_CUBIC)
return res
def extract_square(img, x, y, x0, y0, w0, h0):
x_base = int(x0 + x * w0)
y_base = int(y0 + y * h0)
h0 = int(h0)
w0 = int(w0)
return img[y_base + 12:y_base + h0 - 12, x_base + 10:x_base + w0 - 14]
get_sq_function = find_fit_square
def mark_square(img, x, y, x0, y0, w0, h0, text='X'):
COLOR = 255
if 3 == len(img.shape):
COLOR = (0, 0, 255)
cv2.putText(img, str(text),
(int(x0 + x * w0 + 0.3 * w0), int(y0 + y * h0 + 0.7 * h0)),
cv2.FONT_HERSHEY_SIMPLEX,
1.4, COLOR, 2, cv2.LINE_AA)
cv2.putText(img, "%s" % (x, ),
(int(x0 + x * w0 + 3), int(y0 + y * h0 + 0.6 * h0)),
cv2.FONT_HERSHEY_SIMPLEX,
0.4, COLOR, 1, cv2.LINE_AA)
cv2.putText(img, "%s" % (y, ),
(int(x0 + x * w0 + 3), int(y0 + y * h0 + 0.8 * h0)),
cv2.FONT_HERSHEY_SIMPLEX,
0.4, COLOR, 1, cv2.LINE_AA)
if "__main__" == __name__:
ar = main("hlavolam.jpg")
print(ar)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment