From 15325e0e8faa82f78ab31f7336314799f4c12056 Mon Sep 17 00:00:00 2001
From: linarphy <linarphy@linarphy.net>
Date: Mon, 15 May 2023 16:43:59 +0200
Subject: [PATCH] Add intensity step correction

- Finished curve correction
- Add intensity correction for each x
---
 ETA.py | 514 +++++++++++++++++++++++++++++++--------------------------
 1 file changed, 281 insertions(+), 233 deletions(-)

diff --git a/ETA.py b/ETA.py
index 5823bc6..b17b9e7 100644
--- a/ETA.py
+++ b/ETA.py
@@ -1,6 +1,9 @@
 import numpy as np
+import matplotlib.pyplot as plt
 import utils
 import sys
+import pathlib
+import shelve
 from scipy.signal import convolve as sp_convolve
 from scipy.signal import find_peaks
 from scipy.ndimage import rotate
@@ -9,156 +12,154 @@ if len( sys.argv ) < 2:
     raise Exception( 'this command must have a filename of an ETA fits as an argument' )
 data = utils.load( sys.argv[1] )
 
-"""
-find fill points
-"""
-points = []
+cache_file = pathlib.Path( 'asset/points_' + sys.argv[1].split( '/' )[-1][:-5] + '.pag' )
 
-points += utils.find_point( data[ : , 0 ] , 0 ) # x_min
-points += utils.find_point(
-        data[ : , data.shape[1] - 1 ],
-        data.shape[1] - 1            ,
-) # x_max
-
-index_min = 0
-while data.shape[0] - 1 > index_min:
-    index_min += 1
-    if len( utils.find_point( data[ index_min , : ] , index_min , 'y' ) ) == 3:
-        break
-
-points.append(
-    utils.find_point( data[ index_min , : ] , index_min , 'y' )[1]
-) # y_min
-
-index_max = data.shape[0] - 1
-while index_min < index_max:
-    index_max -= 1
-    if len( utils.find_point( data[ index_max , : ] , index_max , 'y' ) ) == 3:
-        break
-
-
-points.append(
-    utils.find_point( data[ index_max , : ] , index_max , 'y' )[1]
-) # y_max
-
-small_data = utils.compress( data , 5 )
-
-points = utils.big_to_small( points , 5 )
-
-# size - 1
-points[ 1 ][ 1 ] -= 1
-points[ 3 ][ 0 ] -= 1
-
-# little shift to be inside the light
-points[ 2 ][ 1 ] += 3
-points[ 3 ][ 1 ] += 3
-
-"""
-fill data
-"""
-
-extremum = []
-for point in points:
-    if point[0] < points[2][0]:
-        point[0] = points[2][0]
-    if point[1] < points[0][1]:
-        point[1] = points[0][1]
-    taken_points = utils.small_to_big(
-        np.array( [
-            points[2][0],
-            points[0][1],
-        ] ) + utils.fill(
-            small_data[
-                points[2][0] : points[3][0] + 1,
-                points[0][1] : points[1][1] + 1
-            ],
-            [
-                point[0] - points[2][0],
-                point[1] - points[0][1],
-            ],
-            1000
-        ),
-        5
-    )
-    extremum.append( [
-        np.min( taken_points[ : , 1 ] ),
-        np.max( taken_points[ : , 1 ] ),
-        np.min( taken_points[ : , 0 ] ),
-        np.max( taken_points[ : , 0 ] ),
-    ] )
-
-border = {
-    'x': {
-        'min': points[0][1] + extremum[0][1] + 1,
-        'max': points[0][1] + extremum[1][0]    ,
-    },
-    'y': {
-        'min': points[2][0] + extremum[2][3] + 1,
-        'max': points[2][0] + extremum[3][2]    ,
-    },
-}
-
-"""
-label deletion
-"""
-
-mean_data = np.convolve(
-    np.gradient(
-        np.mean(
-            data[
-                border[ 'y' ][ 'min' ] : border[ 'y' ][ 'max' ],
-                border[ 'x' ][ 'min' ] : border[ 'x' ][ 'max' ]
-            ],
-            axis = 0
-        )
-    ),
-    np.ones(
-        int( 0.01 * (
-            border[ 'x' ][ 'max' ] - border[ 'x' ][ 'min' ]
-        ) )
-    ),
-    'same'
-)
-
-mean_data -= np.min( mean_data )
-mean_data /= np.max( mean_data )
-
-top  = utils.consecutive( np.where( mean_data > 0.75 )[0] )
-down = utils.consecutive( np.where( mean_data < 0.25 )[0] )
-
-size_top  = [ len( list_ ) for list_ in top  ]
-size_down = [ len( list_ ) for list_ in down ]
-
-label_x = {
-    'min': border[ 'x' ][ 'min' ] + top[ np.argmax( size_top ) ][0]   ,
-    'max': border[ 'x' ][ 'min' ] + down[ np.argmax( size_down ) ][-1]
-}
-
-if label_x[ 'min' ] < data.shape[1] // 2:
-    if label_x[ 'max' ] < data.shape[1] // 2:
-        border[ 'x' ][ 'min' ] = label_x[ 'max' ]
-    else:
-        raise Exception( 'the label seems to be in the middle of the picture' )
-elif label_x[ 'max' ] > data.shape[1] // 2:
-    border[ 'x' ][ 'max' ] = label_x[ 'min' ]
+if cache_file.is_file():
+    with shelve.open( str( cache_file ) ) as cache:
+        data         = cache[ 'rotated_data' ]
+        border       = cache[ 'border' ]
+        calibrations = cache[ 'calibrations' ]
 else:
-    raise Exception( 'for an unkown reason, label_x[ \'min\' ] > label_x[ \'max\' ]' )
+    """
+    find fill points
+    """
+    points = []
 
-"""
-Rotation
-"""
+    points += utils.find_point( data[ : , 0 ] , 0 ) # x_min
+    points += utils.find_point(
+            data[ : , data.shape[1] - 1 ],
+            data.shape[1] - 1            ,
+    ) # x_max
 
-index    = border[ 'x' ][ 'min' ]
-gradient = np.gradient(
-    data[
-        border[ 'y' ][ 'min' ] : border[ 'y' ][ 'min' ] + (
-            border[ 'y' ][ 'max' ] - border[ 'y' ][ 'min' ]
-        ) // 2,
-        index
-    ]
-)
-while np.max( gradient ) - np.min( gradient ) > 5500:
-    index   += 1
+    index_min = 0
+    while data.shape[0] - 1 > index_min:
+        index_min += 1
+        if len( utils.find_point( data[ index_min , : ] , index_min , 'y' ) ) == 3:
+            break
+
+    points.append(
+        utils.find_point( data[ index_min , : ] , index_min , 'y' )[1]
+    ) # y_min
+
+    index_max = data.shape[0] - 1
+    while index_min < index_max:
+        index_max -= 1
+        if len( utils.find_point( data[ index_max , : ] , index_max , 'y' ) ) == 3:
+            break
+
+
+    points.append(
+        utils.find_point( data[ index_max , : ] , index_max , 'y' )[1]
+    ) # y_max
+
+    small_data = utils.compress( data , 5 )
+
+    points = utils.big_to_small( points , 5 )
+
+    # size - 1
+    points[ 1 ][ 1 ] -= 1
+    points[ 3 ][ 0 ] -= 1
+
+    # little shift to be inside the light
+    points[ 2 ][ 1 ] += 3
+    points[ 3 ][ 1 ] += 3
+
+    """
+    fill data
+    """
+
+    extremum = []
+    for point in points:
+        if point[0] < points[2][0]:
+            point[0] = points[2][0]
+        if point[1] < points[0][1]:
+            point[1] = points[0][1]
+        taken_points = utils.small_to_big(
+            np.array( [
+                points[2][0],
+                points[0][1],
+            ] ) + utils.fill(
+                small_data[
+                    points[2][0] : points[3][0] + 1,
+                    points[0][1] : points[1][1] + 1
+                ],
+                [
+                    point[0] - points[2][0],
+                    point[1] - points[0][1],
+                ],
+                1000
+            ),
+            5
+        )
+        extremum.append( [
+            np.min( taken_points[ : , 1 ] ),
+            np.max( taken_points[ : , 1 ] ),
+            np.min( taken_points[ : , 0 ] ),
+            np.max( taken_points[ : , 0 ] ),
+        ] )
+
+    border = {
+        'x': {
+            'min': points[0][1] + extremum[0][1] + 1,
+            'max': points[0][1] + extremum[1][0]    ,
+        },
+        'y': {
+            'min': points[2][0] + extremum[2][3] + 1,
+            'max': points[2][0] + extremum[3][2]    ,
+        },
+    }
+
+    """
+    label deletion
+    """
+
+    mean_data = np.convolve(
+        np.gradient(
+            np.mean(
+                data[
+                    border[ 'y' ][ 'min' ] : border[ 'y' ][ 'max' ],
+                    border[ 'x' ][ 'min' ] : border[ 'x' ][ 'max' ]
+                ],
+                axis = 0
+            )
+        ),
+        np.ones(
+            int( 0.01 * (
+                border[ 'x' ][ 'max' ] - border[ 'x' ][ 'min' ]
+            ) )
+        ),
+        'same'
+    )
+
+    mean_data -= np.min( mean_data )
+    mean_data /= np.max( mean_data )
+
+    top  = utils.consecutive( np.where( mean_data > 0.75 )[0] )
+    down = utils.consecutive( np.where( mean_data < 0.25 )[0] )
+
+    size_top  = [ len( list_ ) for list_ in top  ]
+    size_down = [ len( list_ ) for list_ in down ]
+
+    label_x = {
+        'min': border[ 'x' ][ 'min' ] + top[ np.argmax( size_top ) ][0]   ,
+        'max': border[ 'x' ][ 'min' ] + down[ np.argmax( size_down ) ][-1]
+    }
+
+    if label_x[ 'min' ] < data.shape[1] // 2:
+        if label_x[ 'max' ] < data.shape[1] // 2:
+            border[ 'x' ][ 'min' ] = label_x[ 'max' ]
+        else:
+            raise Exception( 'the label seems to be in the middle of the picture' )
+    elif label_x[ 'max' ] > data.shape[1] // 2:
+        border[ 'x' ][ 'max' ] = label_x[ 'min' ]
+    else:
+        raise Exception( 'for an unkown reason, label_x[ \'min\' ] > label_x[ \'max\' ]' )
+
+    """
+    Rotation
+    """
+
+    index    = border[ 'x' ][ 'min' ]
     gradient = np.gradient(
         data[
             border[ 'y' ][ 'min' ] : border[ 'y' ][ 'min' ] + (
@@ -167,122 +168,169 @@ while np.max( gradient ) - np.min( gradient ) > 5500:
             index
         ]
     )
-
-positions = np.argmax(
-    sp_convolve(
-        np.gradient(
+    while np.max( gradient ) - np.min( gradient ) > 5500:
+        index   += 1
+        gradient = np.gradient(
             data[
                 border[ 'y' ][ 'min' ] : border[ 'y' ][ 'min' ] + (
                     border[ 'y' ][ 'max' ] - border[ 'y' ][ 'min' ]
-                ) // 2                        ,
-                border[ 'x' ][ 'min' ] : index
-            ]       ,
-            axis = 0
-        )                     ,
-        np.ones( ( 100 , 1 ) ),
-        'valid'
-    )       ,
-    axis = 0
-)
+                ) // 2,
+                index
+            ]
+        )
 
-list_   = np.arange(  0     , index - border[ 'x' ][ 'min' ] , 1 )
-polyval = np.polyfit( list_ , positions                      , 1 )
+    positions = np.argmax(
+        sp_convolve(
+            np.gradient(
+                data[
+                    border[ 'y' ][ 'min' ] : border[ 'y' ][ 'min' ] + (
+                        border[ 'y' ][ 'max' ] - border[ 'y' ][ 'min' ]
+                    ) // 2                        ,
+                    border[ 'x' ][ 'min' ] : index
+                ]       ,
+                axis = 0
+            )                     ,
+            np.ones( ( 100 , 1 ) ),
+            'valid'
+        )       ,
+        axis = 0
+    )
 
-angle = np.arctan( polyval[0] )
+    list_   = np.arange(  0     , index - border[ 'x' ][ 'min' ] , 1 )
+    polyval = np.polyfit( list_ , positions                      , 1 )
 
-data = rotate( data , angle * ( 180 / np.pi ) ) # utils.rotate does not keep intensity absolute value ? TODO
+    angle = np.arctan( polyval[0] )
 
-diff_y = int( np.tan( angle ) * ( border[ 'x' ][ 'max' ] - border[ 'x' ][ 'min' ] ) )
+    data = rotate( data , angle * ( 180 / np.pi ) ) # utils.rotate does not keep intensity absolute value ? TODO
 
-border[ 'y' ][ 'min' ] -= diff_y
-border[ 'y' ][ 'max' ] -= diff_y
+    diff_y = int( np.tan( angle ) * ( border[ 'x' ][ 'max' ] - border[ 'x' ][ 'min' ] ) )
 
-"""
-Calibration
-"""
+    border[ 'y' ][ 'min' ] -= diff_y
+    border[ 'y' ][ 'max' ] -= diff_y
 
-tot_avg = np.mean(
-    data[
-        border[ 'y' ][ 'min' ] : border[ 'y' ][ 'max' ],
-        border[ 'x' ][ 'min' ] : border[ 'x' ][ 'max' ]
-    ]
-)
-def indicator( list_ ):
-    if np.mean( list_ ) > 0.75 * tot_avg:
-        return 0
-    if np.mean( list_ ) < 0.25 * tot_avg:
-        return 1
-    list_ -= np.min( list_ )
-    list_ /= np.max( list_ )
-    positions = np.where( list_ > 0.5 )[0]
-    if len( positions ) < 10:
-        return 2
-    if len( positions ) > 400:
-        return 3
-    distance = np.mean( positions[ 1 : ] - positions[ : -1 ] )
-    if distance < 10:
-        return 4
-    return 10
+    """
+    Calibration
+    """
 
-indicators = np.array( [ indicator( data[ i , border[ 'x' ][ 'min' ] : border[ 'x' ][ 'max' ] ] ) for i in range( border[ 'y' ][ 'min' ] , border[ 'y' ][ 'max' ] , 1 ) ] )
+    tot_avg = np.mean(
+        data[
+            border[ 'y' ][ 'min' ] : border[ 'y' ][ 'max' ],
+            border[ 'x' ][ 'min' ] : border[ 'x' ][ 'max' ]
+        ]
+    )
+    def indicator( list_ ):
+        if np.mean( list_ ) > 0.75 * tot_avg:
+            return 0
+        if np.mean( list_ ) < 0.25 * tot_avg:
+            return 1
+        list_ -= np.min( list_ )
+        list_ /= np.max( list_ )
+        positions = np.where( list_ > 0.5 )[0]
+        if len( positions ) < 10:
+            return 2
+        if len( positions ) > 400:
+            return 3
+        distance = np.mean( positions[ 1 : ] - positions[ : -1 ] )
+        if distance < 10:
+            return 4
+        return 10
 
-calibration_areas = utils.consecutive( np.where( indicators == 10 )[0] )
-calibration_sizes = [ len( calibration_area ) for calibration_area in calibration_areas ]
+    indicators = np.array( [ indicator( data[ i , border[ 'x' ][ 'min' ] : border[ 'x' ][ 'max' ] ].copy() ) for i in range( border[ 'y' ][ 'min' ] , border[ 'y' ][ 'max' ] , 1 ) ] )
 
-y_calibrations = [ calibration_areas[ i ] for i in np.argsort( calibration_sizes ) ][ -2 : ]
-calibrations = {
-    'top': {
-        'x': {
-            'min': border['x']['min'],
-            'max': border['x']['max'],
+    calibration_areas = utils.consecutive( np.where( indicators == 10 )[0] )
+    calibration_sizes = [ len( calibration_area ) for calibration_area in calibration_areas ]
+
+    y_calibrations = [ calibration_areas[ i ] for i in np.argsort( calibration_sizes ) ][ -2 : ]
+    calibrations = {
+        'top': {
+            'x': {
+                'min': border['x']['min'],
+                'max': border['x']['max'],
+            },
+            'y': {
+                'min': border['y']['min'] + y_calibrations[0][ 0],
+                'max': border['y']['min'] + y_calibrations[0][-1],
+            },
         },
-        'y': {
-            'min': border['y']['min'] + y_calibrations[0][ 0],
-            'max': border['y']['min'] + y_calibrations[0][-1],
+        'down': {
+            'x': {
+                'min': border['x']['min'],
+                'max': border['x']['max'],
+            },
+            'y': {
+                'min': border['y']['min'] + y_calibrations[1][ 0],
+                'max': border['y']['min'] + y_calibrations[1][-1],
+            },
         },
-    },
-    'down': {
-        'x': {
-            'min': border['x']['min'],
-            'max': border['x']['max'],
-        },
-        'y': {
-            'min': border['y']['min'] + y_calibrations[1][ 0],
-            'max': border['y']['min'] + y_calibrations[1][-1],
-        },
-    },
-}
+    }
+
+    with shelve.open( str( cache_file ) ) as cache:
+        cache[ 'rotated_data' ] = data
+        cache[ 'border' ]       = border
+        cache[ 'calibrations']  = calibrations
 
 """
 stripes curves detection
 """
-
 list_ = data[
     calibrations[ 'top' ][ 'y' ][ 'max' ] : calibrations[ 'down' ][ 'y' ][ 'min' ],
     border[ 'x' ][ 'min' ] : border[ 'x' ][ 'max' ]
 ].copy()
 list_ -= np.min( list_ , axis = 0 )
 list_ /= np.max( list_ , axis = 0 )
+size     = list_.shape[1]
 
-size  = border[ 'x' ][ 'max' ] - border[ 'x' ][ 'min' ]
-x_stripe = np.arange( border[ 'x' ][ 'min' ] + 1 * size // 4 , border[ 'x' ][ 'min' ] + 3 * size // 4 , 1 ).astype( int )
-y_stripe = np.array( [
-    np.where(
-        list_[ : , x ] > 0.8
-    )[0][0] for x in x_stripe
-] )
+y_stripe = np.argmax( list_ , axis = 0 )
+good_x   = np.where( y_stripe < 2 * np.mean( y_stripe ) )[0]
+x_stripe = np.arange( 0 , size , 1 ).astype( int )[ good_x ]
+y_stripe = y_stripe[ good_x ]
 
 stripes = [ # list of polyval result for each stripe
-    np.polyfit( x_stripe , y_stripe , 2 )
+    np.polyfit( x_stripe , y_stripe , 3 )
 ]
 
 # First deformation
-y_diff  = np.polyval( stripes[0] , np.arange( 0 , size , 1 ) ).astype( int )
-results = np.zeros( ( list_.shape[1] , list_.shape[0] - np.max( y_diff ) ) )
-for i in range( list_.shape[1] ):
-    results[i] = list_[ y_diff[ i ] : list_.shape[0] + y_diff[ i ] - np.max( y_diff ) , i ]
-results = results.transpose()
 
-import matplotlib.pyplot as plt
-plt.imshow( results )
-plt.savefig( 'asset/deformation.png' )
+y_diff  = ( np.polyval( stripes[0] , np.arange( 0 , size , 1 ) ) ).astype( int )
+y_diff[ np.where( y_diff < 0 ) ] = 0
+results = np.zeros( ( list_.shape[0] + np.max( y_diff ) , list_.shape[1] ) )
+for i in range( list_.shape[1] ):
+    results[ : , i ] = np.concatenate( ( np.zeros( np.max( y_diff ) - y_diff[ i ] ) , list_[ : , i ] , np.zeros( y_diff[i] ) ) )
+
+list_results = np.convolve(
+    np.gradient(
+        np.mean( results , axis = 1 ),
+    )            ,
+    np.ones( 50 ),
+    'same'       ,
+)
+
+fall = utils.consecutive( np.where( list_results < - 0.02 )[0] )
+fall = np.array( [
+        np.argmax( list_results )
+    ] + [
+    consecutive[0] + np.argmin(
+        list_results[ consecutive[0] : consecutive[-1] ]
+    ) for consecutive in fall
+] ).astype( int )
+
+"""
+plt.imshow( results , aspect = 'auto' )
+plt.hlines( fall , 0 , size )
+plt.show()
+"""
+
+temp = np.convolve( results[ : , 10000 ] , np.ones( 50 ) , 'same' )
+for i in range( len( fall ) - 1 ):
+    temp[ fall[ i ] : fall[ i + 1 ] ] = np.mean( temp[ fall[ i ] : fall[ i + 1 ] ] )
+
+plt.plot( temp )
+plt.plot(
+    np.convolve(
+        results[ : , 10000 ],
+        np.ones( 50 )       ,
+        'same'              ,
+    ),
+)
+plt.vlines( fall , 0 , 50 , colors = 'red' )
+plt.show()