278 lines
9.8 KiB (Stored with Git LFS)
Python
278 lines
9.8 KiB (Stored with Git LFS)
Python
import cv2
|
|
import argparse
|
|
import enum
|
|
import abc
|
|
from typing import Optional
|
|
import numpy as np
|
|
|
|
class Board_Sizes_ABC(abc.ABC):
|
|
|
|
@abc.abstractmethod
|
|
def get_cell_rel_pos(self, n: int):
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
def get_num_cells(self) -> tuple[int, int]:
|
|
pass
|
|
|
|
def get_cell_position_on_line(self, n: int, p0: np.ndarray, p1: np.ndarray):
|
|
vec = p1 - p0
|
|
return np.round(p0 + vec*self.get_cell_rel_pos(n)).astype(np.int32)
|
|
|
|
def iterat_pos(self) -> tuple[tuple[float, float], tuple[int, int]]:
|
|
for x in range(self.N + 1):
|
|
for y in range(self.N + 1):
|
|
yield (x, y), (self.get_cell_rel_pos(x), self.get_cell_rel_pos(y))
|
|
|
|
class Board_19x19(Board_Sizes_ABC):
|
|
LINE_WIDTH = 1
|
|
CELL_SIZE = 21
|
|
PADDING = 11.5
|
|
N = 18
|
|
|
|
TOTAL_WIDTH = N*CELL_SIZE + (N+1)*LINE_WIDTH + 2*PADDING
|
|
|
|
def __init__(self):
|
|
self.rel_lut = [(self.PADDING + self.LINE_WIDTH*0.5 + (self.LINE_WIDTH + self.CELL_SIZE)*i)/self.TOTAL_WIDTH for i in range(19)]
|
|
|
|
def get_num_cells(self):
|
|
return (self.N+1, self.N+1)
|
|
|
|
def get_cell_rel_pos(self, n: int):
|
|
return self.rel_lut[n]
|
|
|
|
class VISION_PREPROCESSING_MODE(enum.Enum):
|
|
CANNY = "CANNY"
|
|
THRES_HOLD = "THRES_HOLD"
|
|
THRES = "THRES"
|
|
TEST_CANNY_HOLD = "TEST_CANNY_HOLD"
|
|
TEST_CANNY = "TEST_CANNY"
|
|
|
|
class Board_Detection_SM:
|
|
|
|
def __init__(self):
|
|
self.state = VISION_PREPROCESSING_MODE.CANNY
|
|
self.hold_time = 0
|
|
self.dection_history = []
|
|
self.jitter = 0
|
|
self.num_decetions = 0
|
|
self.num_double_decetions = 0
|
|
self.last_decetion_time = 0
|
|
self.last_double_detection_time = 0
|
|
|
|
def advance_state(self, detections: list[np.array]) -> Optional[np.array]:
|
|
detection = self.process_nun_decetions(detections)
|
|
|
|
match self.state:
|
|
case VISION_PREPROCESSING_MODE.CANNY:
|
|
if self.last_decetion_time > 50:
|
|
self.state = VISION_PREPROCESSING_MODE.THRES_HOLD
|
|
self.hold_time = 100
|
|
|
|
case VISION_PREPROCESSING_MODE.THRES_HOLD:
|
|
if self.hold_time == 0:
|
|
self.state = VISION_PREPROCESSING_MODE.THRES
|
|
else:
|
|
self.hold_time -= 1
|
|
|
|
case VISION_PREPROCESSING_MODE.THRES:
|
|
if self.last_decetion_time > 50 or self.jitter > 3 or self.num_double_decetions > 15:
|
|
self.state = VISION_PREPROCESSING_MODE.TEST_CANNY_HOLD
|
|
self.hold_time = 200
|
|
|
|
case VISION_PREPROCESSING_MODE.TEST_CANNY_HOLD:
|
|
if self.hold_time == 0:
|
|
self.state = VISION_PREPROCESSING_MODE.TEST_CANNY
|
|
else:
|
|
self.hold_time -= 1
|
|
|
|
case VISION_PREPROCESSING_MODE.TEST_CANNY:
|
|
if self.last_decetion_time < 10 and self.jitter < 3:
|
|
self.state = VISION_PREPROCESSING_MODE.CANNY
|
|
else:
|
|
self.state = VISION_PREPROCESSING_MODE.THRES_HOLD
|
|
self.hold_time = 100
|
|
|
|
case _:
|
|
raise ValueError("Invalid preprocessing mode")
|
|
|
|
return detection
|
|
|
|
def process_nun_decetions(self, detections: list[np.array]) -> Optional[np.array]:
|
|
num_decetions = len(detections)
|
|
dection = None
|
|
|
|
if num_decetions > 2:
|
|
self.num_double_decetions += 1
|
|
self.num_last_double_detection_time = 0
|
|
else:
|
|
if self.last_double_detection_time > 300:
|
|
self.num_double_decetions = 0
|
|
else:
|
|
self.last_double_detection_time = max(1000, self.last_double_detection_time)
|
|
|
|
if num_decetions == 1:
|
|
self.last_decetion_time = 0
|
|
|
|
self.dection_history.append(np.squeeze(detections[0], axis=1))
|
|
self.dection_history = self.dection_history[-10:]
|
|
else:
|
|
self.last_decetion_time += 1
|
|
|
|
if len(self.dection_history) > 0:
|
|
dection_history_np = np.array(self.dection_history)
|
|
avrage_detection = np.mean(dection_history_np, axis=0)
|
|
dection = avrage_detection.copy()
|
|
|
|
if len(self.dection_history) > 6:
|
|
self.jitter = 0
|
|
for i in dection_history_np:
|
|
self.jitter = np.max(np.abs(i - avrage_detection))
|
|
|
|
return dection
|
|
|
|
def create_info_text(self) -> list[str]:
|
|
|
|
return [
|
|
f"Preprocessing: {self.state.value}",
|
|
f"Double Dection: {self.num_double_decetions}",
|
|
f"Jitter: {self.jitter:.2f}",
|
|
f"Last Dection Time: {self.last_decetion_time}",
|
|
f"Last Double Dection Time: {self.last_double_detection_time}",
|
|
f"Hold Time: {self.hold_time}"
|
|
]
|
|
|
|
|
|
|
|
|
|
def detect_board(state: VISION_PREPROCESSING_MODE, frame: np.ndarray) -> tuple[list[np.array], np.ndarray]:
|
|
|
|
frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
|
|
|
if state == VISION_PREPROCESSING_MODE.CANNY or \
|
|
state == VISION_PREPROCESSING_MODE.TEST_CANNY or \
|
|
state == VISION_PREPROCESSING_MODE.TEST_CANNY_HOLD:
|
|
|
|
frame_proc = cv2.GaussianBlur(frame_gray, (3, 3), 0)
|
|
frame_proc = cv2.Canny(frame_proc, 130, 200)
|
|
|
|
elif state == VISION_PREPROCESSING_MODE.THRES_HOLD or state == VISION_PREPROCESSING_MODE.THRES:
|
|
frame_proc = cv2.GaussianBlur(frame_gray, (3, 3), 0)
|
|
frame_proc = cv2.adaptiveThreshold(frame_proc, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 19, 3)
|
|
|
|
else:
|
|
raise ValueError("Invalid preprocessing mode")
|
|
|
|
countours, _ = cv2.findContours(frame_proc, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
return list(
|
|
filter(lambda x: type(x) == np.ndarray,
|
|
map(filter_countours, countours)
|
|
)
|
|
), frame_proc
|
|
|
|
def filter_countours(c: np.array) -> Optional[np.array]:
|
|
|
|
esp = cv2.arcLength(c, True)*0.05
|
|
approx = cv2.approxPolyDP(c, esp, True)
|
|
|
|
# Stage 1: Has 4 corners
|
|
if len(approx) != 4:
|
|
return None
|
|
|
|
# Stage 2: Area is big enough
|
|
area = cv2.contourArea(approx)
|
|
|
|
if area < 1000:
|
|
return None
|
|
|
|
# Stage 3: Are all side findContourslengths similar
|
|
lengths = np.array([
|
|
np.linalg.norm(approx[0] - approx[1]),
|
|
np.linalg.norm(approx[1] - approx[2]),
|
|
np.linalg.norm(approx[2] - approx[3]),
|
|
np.linalg.norm(approx[3] - approx[0])
|
|
])
|
|
|
|
avrg_length = np.mean(lengths)
|
|
lengths_diff = np.abs(lengths)/avrg_length
|
|
|
|
if not (np.all(lengths_diff < 1.1) and np.all(lengths_diff > 0.90)):
|
|
return None
|
|
|
|
return approx
|
|
|
|
def draw_board_markers(f: np.array, conture: np.array, board_messurements: Board_Sizes_ABC):
|
|
N_line_positions = []
|
|
S_line_positions = []
|
|
W_line_positions = []
|
|
E_line_positions = []
|
|
|
|
for i in range(19):
|
|
N_line_positions.append(board_messurements.get_cell_position_on_line(i, conture[0], conture[1]))
|
|
E_line_positions.append(board_messurements.get_cell_position_on_line(i, conture[1], conture[2]))
|
|
S_line_positions.append(board_messurements.get_cell_position_on_line(i, conture[2], conture[3]))
|
|
W_line_positions.append(board_messurements.get_cell_position_on_line(i, conture[3], conture[0]))
|
|
|
|
N_line_positions = np.array(N_line_positions, dtype=np.int32)
|
|
E_line_positions = np.array(E_line_positions, dtype=np.int32)
|
|
S_line_positions = np.array(list(reversed(S_line_positions)), dtype=np.int32)
|
|
W_line_positions = np.array(list(reversed(W_line_positions)), dtype=np.int32)
|
|
|
|
cv2.circle(f, conture[0].astype(np.int32), 3, (0, 255, 255), -1)
|
|
cv2.circle(f, conture[1].astype(np.int32), 3, (0, 255, 255), -1)
|
|
cv2.circle(f, conture[2].astype(np.int32), 3, (0, 255, 255), -1)
|
|
cv2.circle(f, conture[3].astype(np.int32), 3, (0, 255, 255), -1)
|
|
|
|
for p in N_line_positions:
|
|
cv2.circle(f, p, 2, (0, 0, 255), -1)
|
|
for p in E_line_positions:
|
|
cv2.circle(f, p, 2, (0, 0, 255), -1)
|
|
for p in S_line_positions:
|
|
cv2.circle(f, p, 2, (0, 0, 255), -1)
|
|
for p in W_line_positions:
|
|
cv2.circle(f, p, 2, (0, 0, 255), -1)
|
|
|
|
for p0, p1 in zip(N_line_positions, S_line_positions):
|
|
cv2.line(f, tuple(p0), tuple(p1), (0, 0, 255), 1)
|
|
|
|
for i in range(19):
|
|
p2 = board_messurements.get_cell_position_on_line(i, p0, p1)
|
|
cv2.circle(f, p2, 2, (0, 0, 255), -1)
|
|
|
|
for p0, p1 in zip(E_line_positions, W_line_positions):
|
|
cv2.line(f, tuple(p0), tuple(p1), (0, 255, 0), 1)
|
|
|
|
for i in range(19):
|
|
p2 = board_messurements.get_cell_position_on_line(i, p0, p1)
|
|
cv2.circle(f, p2, 2, (0, 255, 0), -1)
|
|
|
|
cv2.drawContours(f, [conture], -1, (0, 255, 255), 1)
|
|
|
|
def corret_board_perspective(frame: np.array, conture: np.array, s=250):
|
|
p1 = conture.astype(np.float32)
|
|
p2 = np.array([
|
|
[s,0], # Top right
|
|
[0, 0], # Top Left
|
|
[0,s], # Bottom Right
|
|
[s,s], # Bottom Left
|
|
], dtype=np.float32)
|
|
|
|
T = cv2.getPerspectiveTransform(p1, p2)
|
|
res = cv2.warpPerspective(frame, T, (s, s))
|
|
res = res[0:s,0:s]
|
|
|
|
return res
|
|
|
|
def pad_center(frame: np.array, target_size: np.array):
|
|
f = frame[0:target_size[1],0:target_size[0]]
|
|
|
|
horiztontal_padding = (target_size[1]-f.shape[1]) // 2
|
|
vertical_padding = (target_size[0]-f.shape[0]) // 2
|
|
|
|
return cv2.copyMakeBorder(
|
|
f,
|
|
vertical_padding, vertical_padding,
|
|
horiztontal_padding, horiztontal_padding,
|
|
cv2.BORDER_CONSTANT, value=(0, 0, 0)
|
|
)
|
|
|