From 8db571f34cbc49c1194adfbdfb123387b79ecb93 Mon Sep 17 00:00:00 2001
From: linarphy <linarphy@linarphy.net>
Date: Mon, 22 May 2023 16:01:33 +0200
Subject: [PATCH] Add retrieve_peaks() function and fix consecutive()

---
 utils.py | 44 +++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 43 insertions(+), 1 deletion(-)

diff --git a/utils.py b/utils.py
index bf0a33f..514b3bc 100644
--- a/utils.py
+++ b/utils.py
@@ -192,7 +192,7 @@ def last_consecutive( list_ ):
         raise ValueError( 'list_ must be a list, ' + type( list_ ) + ' given' )
     first , lower , greater = list_[0] , 0 , len( list_ )
     i = lower + ( greater - lower ) // 2
-    while greater - lower != 1:
+    while greater - lower != 0:
         i = lower + ( greater - lower ) // 2
         if list_[ i ] - first != i: # outside of the consecutive list
             greater = i
@@ -212,3 +212,45 @@ def rotate( image , angle ):
     cX , cY        = ( width // 2 , height // 2 )
     matrix         = cv2.getRotationMatrix2D( ( cX , cY ) , angle , 1 )
     return cv2.warpAffine( image , matrix , ( width , height ) , flags = cv2.INTER_NEAREST )
+def retrieve_peaks( data , window_size = 5 , error_coef = 1.05 , max_window_size = 30 , min_successive = 3 ):
+    """
+    get peak position from a 1D data
+    """
+    spectral_energy    = np.log( data ** 2 )
+    error_thr          = error_coef / np.median( spectral_energy )
+
+    average_window = np.convolve(
+        spectral_energy       ,
+        np.ones( window_size ),
+        'same'                ,
+    ) / window_size
+    average_energy = np.mean( average_window )
+    peaks          = np.where(
+        average_window / average_energy ** 2 > error_thr
+    )[0]
+    peaks          = [
+        np.mean( peak ) for peak in consecutive( peaks )
+    ]
+    successive = 0
+
+    while successive < min_successive and window_size < max_window_size:
+        average_window = np.convolve(
+            spectral_energy       ,
+            np.ones( window_size ),
+            'same'                ,
+        ) / window_size
+        average_energy = np.mean( average_window )
+        new_peaks      = np.where(
+            average_window / average_energy ** 2 > error_thr
+        )[0]
+        new_peaks = [
+            np.mean( peak ) for peak in consecutive( new_peaks )
+        ]
+
+        if len( peaks ) == len( new_peaks ):
+            successive += 1
+        else:
+            successive = 0
+            peaks      = new_peaks
+        window_size += 1
+    return peaks