Skip to content

Commit

Permalink
refactoring, auto-scale images in "display_imgs"
Browse files Browse the repository at this point in the history
  • Loading branch information
sitzmann committed Jan 26, 2022
1 parent 2db5c76 commit 81f0de9
Showing 6 changed files with 34 additions and 23 deletions.
2 changes: 2 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -1,2 +1,4 @@
.idea/
data/MNIST/
input_imgs/
media_out/
9 changes: 8 additions & 1 deletion cfg.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import cv2

# coding=utf-8
# directories
input_imgs_dir = 'input_imgs/'
@@ -11,6 +13,7 @@
resolution_x, resolution_y = 1920, 1080
max_num_windows = 5
digit_target_size = 28
tactic_board_height_cm, tactic_board_width_cm = 100, 37

#
min_intensity, max_intensity, medium_intensity = 0, 255, 100
@@ -22,18 +25,22 @@
ksize_initial_blur = 15
field_detection_poly_epsilon = 150
offset_binarize_global = -3
ksize_thresh_field = 15
offset_thresh_field = 1

# player detection
ksize_blur_crop = 11
filled = -1


radius_players_cm = 1
player_radius_lb, player_radius_ub = 0.75, 0.9
h_circles_args = dict(method=cv2.HOUGH_GRADIENT, dp=1.8, param1=150, param2=30)

# digits
font_size = 0.15 * resize_factor


# drawing
border_size_m = 3
draw_scale = 8
draw_scale = 8.7
11 changes: 7 additions & 4 deletions cv_utils.py
Original file line number Diff line number Diff line change
@@ -30,9 +30,12 @@ def rotate_img(img, angle):
return cv2.warpAffine(img, rot_mat, img.shape[1::-1], flags=cv2.INTER_LINEAR)


def display_img(img, scale=8, window_name=None, wait=True, pos=None):
def display_img(img, window_name=None, wait=True, pos=None, scale=None):
global show_count
scale /= cfg.resize_factor
if scale is None:
scale = cfg.resolution_x / cfg.max_num_windows / img.shape[1]
else:
scale /= cfg.resize_factor
pos = show_count if pos is None else pos
h, w = (np.array(img.shape[:2]) * scale).astype(np.int32)
window_name = f'{show_count}' if window_name is None else window_name
@@ -43,9 +46,9 @@ def display_img(img, scale=8, window_name=None, wait=True, pos=None):
cv2.waitKey(0) if wait else None


def display_imgs(imgs, scale, start_pos=0):
def display_imgs(imgs, start_pos=0, scale=None):
for i, img in enumerate(imgs):
display_img(img, scale, wait=False, pos=i + start_pos)
display_img(img, str(i), False, i + start_pos, scale)
cv2.waitKey(0)


7 changes: 3 additions & 4 deletions drawer.py
Original file line number Diff line number Diff line change
@@ -7,6 +7,7 @@
import shutil
import state
import cfg
import cv_utils


def main():
@@ -81,13 +82,11 @@ def rel_line(x, y):
ctx.rel_line_to(*m2p([x, y]))


def show(surface, filename='temp', wait=10):
def show(surface, filename='temp', wait=10, pos=4):
os.makedirs(cfg.media_out_dir, exist_ok=True)
path = f'{cfg.media_out_dir}/{filename}.png'
surface.write_to_png(path)
cv2.imshow(filename, cv2.imread(path))
window_x = cfg.resolution_x - 50 - m2p(cfg.field_width_m + 2 * cfg.border_size_m)
cv2.moveWindow(filename, window_x, 0)
cv_utils.display_img(cv2.imread(path), filename, False, pos)
cv2.waitKey(wait)


6 changes: 3 additions & 3 deletions main.py
Original file line number Diff line number Diff line change
@@ -5,7 +5,7 @@

def main():
# for img_path in [imgs_dir + 'ho-stack-1.jpg', imgs_dir + 'ho-stack-2.jpg']:
img_path = imgs_dir + 'disc.jpg'
img_path = 'input_imgs/not_supported/arrows_and_areas.jpg'
show_digits = True
show_circles = True
record_examples = False
@@ -24,5 +24,5 @@ def animate():


if __name__ == '__main__':
# main()
animate()
main()
# animate()
22 changes: 11 additions & 11 deletions scan.py
Original file line number Diff line number Diff line change
@@ -36,11 +36,11 @@ def scan(img_path: str, show_digits=False, show_circles=False, labeling_mode=Fal
gray_sharp = cv_utils.min_max_normalize(np.clip(gray - gray_blurred, cfg.min_intensity, cfg.max_intensity))
binary = cv_utils.adaptive_threshold(gray_sharp, ksize, cfg.offset_binarize_global)
binary = cv2.medianBlur(binary, cfg.ksize_blur_thresholded)
radius_pixels = cfg.radius_players_cm * new_h // cfg.field_height_m
bounds = [int(radius_pixels * factor) for factor in [cfg.player_radius_lb, cfg.player_radius_ub]]
circles = cv2.HoughCircles(binary, cv2.HOUGH_GRADIENT, 1.8, new_h/100, None, 150, 30, *bounds)[0].astype(np.uint16)
radius_pixels = cfg.radius_players_cm * new_h // cfg.tactic_board_height_cm
lb, ub = [int(radius_pixels * factor) for factor in [cfg.player_radius_lb, cfg.player_radius_ub]]
circles = cv2.HoughCircles(binary, minDist=new_h/100, minRadius=lb, maxRadius=ub, **cfg.h_circles_args)[0]
players_mask = np.zeros_like(binary)
[cv2.circle(players_mask, (c[0], c[1]), radius_pixels, cfg.max_intensity, -1) for c in circles]
[cv2.circle(players_mask, (c[0], c[1]), radius_pixels, cfg.max_intensity, -1) for c in circles.astype(np.uint16)]
annotated, players = img.copy(), []
for c in cv_utils.find_contours(players_mask):
players.append(identify_player(img.copy(), c, radius_pixels, show_digits, labeling_mode))
@@ -90,7 +90,7 @@ def detect_field(img: np.ndarray, show_edges=False) -> np.ndarray:
"""
img_gray = cv_utils.min_max_normalize(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY))
img_gray = cv2.medianBlur(img_gray, cfg.ksize_initial_blur)
edges = cv_utils.adaptive_threshold(img_gray, 15, 1)
edges = cv_utils.adaptive_threshold(img_gray, cfg.ksize_thresh_field, cfg.offset_thresh_field)
cv_utils.display_img(edges, window_name='edges') if show_edges else None
contours, _ = cv2.findContours(edges, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
corners = max_area = None
@@ -117,15 +117,15 @@ def identify_player(img, contour, radius_pixels, show_digits, labeling_mode=Fals
"""
cropped_contour = np.zeros_like(img[:, :, 0])
pos = cv_utils.get_contour_center(contour)
cv2.drawContours(cropped_contour, [contour], 0, cfg.max_intensity, -1)
cv2.drawContours(cropped_contour, [contour], 0, cfg.max_intensity, cfg.filled)
crop = cv_utils.crop_to_content(cropped_contour, img, radius_pixels // 2)
crop_binary, digit_hull = extract_digit(crop)
angle, frame_contour = estimate_frame(crop, digit_hull, show_digits)
crop_rotated = rotate_to_mnist(crop_binary, -angle)
classification = digit_classification.classify_img_by_examples(crop_rotated, show=show_digits)
if show_digits:
for i, img in enumerate([crop_binary, crop_rotated]):
cv_utils.display_img(img, 70, str(i + 3), False, i + 3)
cv_utils.display_img(img, str(i + 3), False, i + 3)
cv2.waitKey(0)
background_mask = cv_utils.crop_to_content(cropped_contour, cropped_contour, radius_pixels // 2)
for c in [frame_contour, digit_hull]: # digit_contour would be more accurate than digit_hull
@@ -168,7 +168,7 @@ def estimate_frame(crop, digit_hull, show_digits):
cv2.line(frame_2, center.astype(np.uint8), opening_point.astype(np.uint8), cfg.medium_intensity, 1)
if show_digits:
for i, img in enumerate([gray_sharp, frame, frame_2]):
cv_utils.display_img(img, 70, str(i), False, i)
cv_utils.display_img(img, str(i), False, i)
cv2.waitKey(100)
return orientation, frame_contour

@@ -215,7 +215,7 @@ def extract_digit(img):
min_area, best_hull = area, hull
if best_hull is None:
cv2.drawContours(img, contours, -1, (cfg.max_intensity, 100, 0), 1)
cv_utils.display_imgs([img, crop_binary], 100)
cv_utils.display_imgs([img, crop_binary])
digit_hull_mask = np.zeros_like(img[:, :, 0])
cv2.drawContours(digit_hull_mask, [best_hull], 0, cfg.max_intensity, -1)
crop_binary[digit_hull_mask == 0] = 0
@@ -228,7 +228,7 @@ def label_refenrence_player(img):
asks the user to input the player label (1-7) and saves label + image to recognize it in the future
:param img: MNIST-like image
"""
cv_utils.display_img(img, 100, 'digit', wait=False, pos=1)
cv_utils.display_img(img, 'digit', wait=False, pos=1)
cv2.waitKey(100)
label = input('which digit is this?\n')
dirname = f'{cfg.reference_digits_dir}/{label}'
@@ -239,6 +239,6 @@ def label_refenrence_player(img):


if __name__ == '__main__':
state = scan(cfg.input_imgs_dir + 'disc.jpg', show_digits=False, show_circles=True)
state = scan(cfg.input_imgs_dir + 'ho-stack-1.jpg', show_digits=False, show_circles=True)
surface = drawer.draw_scene(state)
drawer.show(surface, wait=0)

0 comments on commit 81f0de9

Please sign in to comment.