Skip to content

Instantly share code, notes, and snippets.

@petrblahos
Created October 15, 2018 18:32
Show Gist options
  • Save petrblahos/6a0527a76847de70481097cec103c717 to your computer and use it in GitHub Desktop.
Save petrblahos/6a0527a76847de70481097cec103c717 to your computer and use it in GitHub Desktop.
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