Add rotation step for spectrum.py

spectrum.py

@@ -1,10 +1,12 @@
 import numpy as np
 import matplotlib.pyplot as plt
-import scipy.signal as signal
+from scipy.signal import medfilt , find_peaks
+from scipy.optimize import curve_fit
 import utils
 import sys
 import shelve
 import pathlib
+from scipy.ndimage import rotate
 
 cache , filename , output , calibration , verbose , no_cache = '' , None , None , None , False, False
 if len( sys.argv ) < 2:
@@ -219,6 +221,73 @@ else:
 
     if verbose:
         print( 'first zoom finished'      )
+        print( 'starting rotation'        )
+        print( 'retrieving current angle' )
+
+    """
+    Rotation
+    """
+
+    gauss = lambda x , sigma , mu , a , b : a * (
+            1 / ( sigma * np.sqrt( 2 * np.pi ) )
+        ) * np.exp(
+            - ( x - mu ) ** 2 / ( 2 * sigma ** 2 )
+        ) + b
+    guess_params = [
+        1                                                      ,
+        ( border[ 'x' ][ 'max' ] - border[ 'x' ][ 'min' ] ) / 2,
+        np.mean( data[
+            border[ 'y' ][ 'min' ] : border[ 'y' ][ 'max' ],
+            border[ 'x' ][ 'min' ] : border[ 'x' ][ 'max' ]
+        ] )                                                    ,
+        np.mean( data[
+            border[ 'y' ][ 'min' ] : border[ 'y' ][ 'max' ],
+            border[ 'x' ][ 'min' ] : border[ 'x' ][ 'max' ]
+        ] )                                                    ,
+    ]
+    number         = 1000
+    position_peaks = np.zeros( number )
+    indexes        = np.linspace( border[ 'x' ][ 'min' ] , border[ 'x' ][ 'max' ] , number , dtype = int )
+    for i in range( number ):
+        x = np.arange(
+            border[ 'y' ][ 'min' ],
+            border[ 'y' ][ 'max' ],
+            1                     ,
+        )
+        y = data[
+            border[ 'y' ][ 'min' ] : border[ 'y' ][ 'max' ],
+            indexes[ i ]
+        ]
+
+        try:
+            position_peaks[ i ] = curve_fit(
+                gauss       ,
+                x           ,
+                y           ,
+                guess_params,
+            )[0][1]
+        except:
+            position_peaks[ i ] = 0
+    position_peaks = medfilt(
+        position_peaks[ np.where( position_peaks != 0 ) ],
+        11                                               ,
+    )
+    abciss         = np.arange(
+        len( position_peaks )
+    )[ np.where( position_peaks ) ]
+    polyval        = np.polyfit( abciss , position_peaks , 1 )
+
+    angle = np.arctan( polyval[0] )
+
+    if verbose:
+        print( 'current angle retrieved: ' + str( angle ) )
+        print( 'starting image rotation'                  )
+
+    data = rotate( data , angle * ( 180 / np.pi ) ) # utils.rotate does not keep intenisty absolute value TODO
+
+    if verbose:
+        print( 'image rotation finished' )
+        print( 'rotation finished'           )
         print( 'starting spectrum isolation' )
         print( 'starting y border detection' )
 
@@ -310,7 +379,7 @@ else:
         list_ /= np.max( list_ ) # max 1
 
         amplitude = np.mean( list_ ) # the lower the better
-        peaks     = signal.find_peaks(
+        peaks     = find_peaks(
             list_                                                   ,
             height = np.mean( list_ ) + ( 1 - np.mean( list_ ) ) / 2,
         )[0]
@@ -405,6 +474,32 @@ Calibration
 if calibration != None:
     if verbose:
         print( 'starting calibration' )
+
+    mean_up = np.mean( data[
+        calibrations[ 'top' ][ 'y' ][ 'min' ] : calibrations[ 'top' ][ 'y' ][ 'max' ],
+        calibrations[ 'top' ][ 'x' ][ 'min' ] : claibrations[ 'top' ][ 'x' ][ 'max' ]
+    ] , axis = 0 )
+    mean_up = np.mean( data[
+        calibrations[ 'down' ][ 'y' ][ 'min' ] : calibrations[ 'down' ][ 'y' ][ 'max' ],
+        calibrations[ 'down' ][ 'x' ][ 'min' ] : claibrations[ 'down' ][ 'x' ][ 'max' ]
+    ] , axis = 0 )
+    peaks_up = np.array(
+        utils.retrieve_peaks(
+            mean_up            ,
+            window_size = 1    ,
+            max_window_size = 1,
+        )          ,
+        dtype = int,
+    )
+    peaks_down = np.array(
+        utils.retrieve_peaks(
+            mean_down          ,
+            window_size = 1    ,
+            max_window_size = 1,
+        )          ,
+        dtype = int,
+    )
+
     if verbose:
         print( 'calibration finished' )