41 lines
1.6 KiB (Stored with Git LFS)
Python
41 lines
1.6 KiB (Stored with Git LFS)
Python
import cv2
|
|
import numpy as np
|
|
|
|
class PREPROCESSING_MODE:
|
|
CANNY = "CANNY"
|
|
THRES = "THRES"
|
|
|
|
class Preprocessing_parameters:
|
|
def __init__(self, low_threshold: int=130,
|
|
high_threshold: int=150,
|
|
kernel_size: int=19,
|
|
threshold: int=3,
|
|
gaussian_kernel: tuple[int, int]=(3, 3)
|
|
) -> None:
|
|
self.canny_low_threshold = low_threshold
|
|
self.canny_high_threshold = high_threshold
|
|
self.thres_kernel_size = kernel_size
|
|
self.thres_threshold = threshold
|
|
self.gaussian_kernel = gaussian_kernel
|
|
|
|
def canny_edge_dection(frame: np.ndarray, parameters: Preprocessing_parameters) -> np.ndarray:
|
|
frame_proc = cv2.GaussianBlur(frame, parameters.gaussian_kernel, 0)
|
|
return cv2.Canny(
|
|
frame_proc,
|
|
parameters.canny_low_threshold, parameters.canny_high_threshold
|
|
)
|
|
|
|
def thresholding_edge_detection(frame: np.ndarray, parameters: Preprocessing_parameters) -> np.ndarray:
|
|
frame_proc = cv2.GaussianBlur(frame, parameters.gaussian_kernel, 0)
|
|
return cv2.adaptiveThreshold(
|
|
frame_proc, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,
|
|
parameters.thres_kernel_size, parameters.thres_threshold
|
|
)
|
|
|
|
def edge_dection(frame: np.ndarray, mode: PREPROCESSING_MODE, param: Preprocessing_parameters) -> np.ndarray:
|
|
if mode == PREPROCESSING_MODE.CANNY:
|
|
return canny_edge_dection(frame, param)
|
|
elif mode == PREPROCESSING_MODE.THRES:
|
|
return thresholding_edge_detection(frame, param)
|
|
else:
|
|
raise ValueError("Invalid mode") |