From d48eb53a70076799cba9dec62d450255ca3ef11c Mon Sep 17 00:00:00 2001 From: raccoonliukai <903896015@qq.com> Date: Wed, 16 Oct 2019 20:20:18 +0800 Subject: [PATCH] fix(preprocessor): add max_shot_num for shotdetect --- gnes/preprocessor/helper.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/gnes/preprocessor/helper.py b/gnes/preprocessor/helper.py index b7be20c8..d1ef4b2e 100644 --- a/gnes/preprocessor/helper.py +++ b/gnes/preprocessor/helper.py @@ -351,13 +351,17 @@ def motion_algo(distances: List[float], **kwargs) -> List[int]: min_dist = kwargs.get('min_dist', 10) motion_step = kwargs.get('motion_step', 15) neigh_avg = kwargs.get('neigh_avg', 2) + max_shot_num = kwargs.get('max_shot_num', 30) - 1 shots = [] num_frames = len(distances) + 2 * neigh_avg + 1 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] - + if len(p) > max_shot_num: + max_distances = np.array(distances)[p] + top = np.argsort(-max_distances)[:max_shot_num] + p = p[np.sort(top)] shots.append(0) shots.append(p[0] + neigh_avg + 1) for i in range(1, len(p)):