Skip to content
Permalink
Browse files

fix(shotdetect): support get arguments from yaml

  • Loading branch information...
raccoonliukai committed Sep 27, 2019
1 parent 17aa78d commit ca73b7020d0bbef1c109b1fe78c95759a3fa258c
Showing with 16 additions and 20 deletions.
  1. +15 −19 gnes/preprocessor/helper.py
  2. +1 −1 gnes/preprocessor/video/shotdetect.py
@@ -215,21 +215,17 @@ def hsv_histogram(image: 'np.ndarray') -> 'np.ndarray':
def canny_edge(image: 'np.ndarray', **kwargs) -> 'np.ndarray':
import cv2

arg_dict = {
'sigma': 0.5,
'gauss_kernel': (9, 9),
'l2_gradient': True
}
arg_dict.update(kwargs)
sigma = kwargs.get('sigma', 0.5)
gauss_kernel = kwargs.get('gauss_kernel', (9, 9))
l2_gradient = kwargs.get('l2_gradient', True)

image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
# apply automatic Canny edge detection using the computed median
v = np.median(image)
sigma = arg_dict['sigma']
low_threshold = ((1.0 - sigma) * v).astype("float32")
high_threshold = ((1.0 + sigma) * v).astype("float32")
tmp_image = cv2.GaussianBlur(image, arg_dict['gauss_kernel'], 1.2)
edge_image = cv2.Canny(tmp_image, low_threshold, high_threshold, L2gradient=arg_dict['l2_gradient'])
tmp_image = cv2.GaussianBlur(image, gauss_kernel, 1.2)
edge_image = cv2.Canny(tmp_image, low_threshold, high_threshold, L2gradient=l2_gradient)
return edge_image


@@ -256,10 +252,12 @@ def compute_descriptor(image: 'np.ndarray',
return funcs[method](image)


def compare_ecr(descriptors: List['np.ndarray'], dilate_rate: int = 5, neigh_avg: int = 2) -> List[float]:
def compare_ecr(descriptors: List['np.ndarray'], **kwargs) -> List[float]:
import cv2

""" Apply the Edge Change Ratio Algorithm"""
dilate_rate = kwargs.get('dilate_rate', 5)
neigh_avg = kwargs.get('neigh_avg', 2)
divd = lambda x, y: 0 if y == 0 else x / y

dicts = []
@@ -348,28 +346,26 @@ def thre_algo(distances: List[float], **kwargs) -> List[int]:

def motion_algo(distances: List[float], **kwargs) -> List[int]:
import peakutils
arg_dict = {
'threshold': 0.6,
'min_dist': 10,
'motion_step': 15
}
arg_dict.update(kwargs)

threshold = kwargs.get('threshold', 0.6)
min_dist = kwargs.get('min_dist', 10)
motion_step = kwargs.get('motion_step', 15)
neigh_avg = kwargs.get('neigh_avg', 2)

shots = []
num_frames = len(distances) + 2 * neigh_avg + 1
p = peakutils.indexes(np.array(distances).astype('float32'), thres=arg_dict['threshold'], min_dist=arg_dict['min_dist']) if len(distances) else []
p = peakutils.indexes(np.array(distances).astype('float32'), thres=threshold, min_dist=min_dist) if len(distances) else []
if len(p) == 0:
return [0, num_frames]

shots.append(0)
shots.append(p[0] + neigh_avg + 1)
for i in range(1, len(p)):
# We check that the peak is not due to a motion in the image
valid_dist = not arg_dict['motion_step'] or not check_motion(distances[p[i]-arg_dict['motion_step']:p[i]], distances[p[i]])
valid_dist = not motion_step or not check_motion(distances[p[i]-motion_step:p[i]], distances[p[i]])
if valid_dist:
shots.append(p[i] + neigh_avg + 1)
if shots[-1] < num_frames - arg_dict['min_dist']:
if shots[-1] < num_frames - min_dist:
shots.append(num_frames)
elif shots[-1] > num_frames:
shots[-1] = num_frames
@@ -57,7 +57,7 @@ def detect_shots(self, frames: 'np.ndarray') -> List[List['np.ndarray']]:

# compute distances between frames
if self.distance_metric == 'edge_change_ration':
dists = compare_ecr(descriptors)
dists = compare_ecr(descriptors, **self._detector_kwargs)
else:
dists = [
compare_descriptor(pair[0], pair[1], self.distance_metric)

0 comments on commit ca73b70

Please sign in to comment.
You can’t perform that action at this time.