fcd/data/data_utils.py at main · i4Ds/fcd · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
import numpy as np
from ..utils.print_utils import p2, p, d2


def sta_info( np_arr, rest_no_show=True ):
    p( '-------- Statistical Info. --------')
    p2( 'Mean', np_arr.mean() )
    p2( 'Min', np_arr.min() )
    p2( 'Max', np_arr.max() )
    p2( 'Median', np.median(np_arr) )
    p2( 'Std-dev', np_arr.std() )

    if not rest_no_show:
        # Calculate Interquartile Range (IQR)
        p2( 'Q1', np.percentile(np_arr, 25) ) # first quartile (Q1)
        p2( 'Q3', np.percentile(np_arr, 75) ) # third quartile (Q3)
        p2( 'IQR', np.percentile(np_arr, 75) - np.percentile(np_arr, 25) )


def combine_data(*data):
    return np.concatenate(data, axis=0)


def combine_npz_data(*npz_data, out_file=None):
    '''
    Combine multiple NPZ files into a single NPZ file.
    Example:
    data_1 = np.load("stx_nos_1.npz")
    data_2 = np.load("stx_nos_2.npz")
    data_3 = np.load("stx_nos_3.npz")
    data_4 = np.load("stx_nos_4.npz")
    data = combine_npz_data(data_1, data_2, data_3, data_4, out_file="stx_labelled_data.npz")
    '''
    # Dictionary to hold combined arrays
    combined_arrays = {}
    i = 0

    for data in npz_data:
        i += 1
        # Iterate through arrays in the current file
        for arr_name in data.files:
            # If the array already exists in our combined dictionary,
            # we concatenate it with the new one. Otherwise, we just add it.
            if arr_name in combined_arrays:
                combined_arrays[arr_name] = np.concatenate((combined_arrays[arr_name], data[arr_name]), axis=0)
            else:
                combined_arrays[arr_name] = data[arr_name]

    p2(f'{i} arrays are combined on', combined_arrays.keys())

    if out_file is not None:
        # Save the combined arrays into a new NPZ file
        np.savez(out_file, **combined_arrays)
        p2('Saved combined arrays to', out_file)

    return combined_arrays


def to_df_row(metric, np_arr):
    return [metric,
            "{:.3f}".format(np_arr.mean()),
            # "{:.3f}".format(np_arr.min()),
            # "{:.3f}".format(np_arr.max()),
            # "{:.3f}".format(np.median(np_arr)),
            "{:.3f}".format(np.std(np_arr)),
            ]


def get_alpha( vis_arr, pix_size ):
    # Compute alpha for a single vis array
    norm_vis = np.sqrt( np.square( vis_arr[:24] ) + np.square( vis_arr[24:] ) )
    alpha = 2 * np.max( norm_vis ) / (pix_size * pix_size)
    return alpha


def norm_couple( vis_arr, gt_arr, pix_size ):
    """
    vis_data: visibility data
    gt_data: data of ground truth
    """
    alpha = get_alpha( vis_arr, pix_size )
    return vis_arr/alpha, gt_arr/alpha, alpha


def norm_data(vis_data, gt_data, pix_size=2.0):
    # Initialize arrays to hold normalised data
    norm_vis_data = np.empty_like(vis_data)
    norm_gt_data = np.empty_like(gt_data)
    alpha_vec = np.empty(len(vis_data), dtype=np.float32)

    # Process each vis and gt pair
    for i in range(len(vis_data)):
        # Normalise visibility and ground truth data
        norm_vis_data[i], norm_gt_data[i], alpha_vec[i] = \
            norm_couple(vis_data[i], gt_data[i], pix_size)

    return norm_vis_data, norm_gt_data, alpha_vec


def generate_sample(vis_data, img_data, size=3):
    indices = np.random.choice(len(vis_data), size, replace=False)
    # indices = [129229, 155731, 57992]
    val_vis = vis_data[indices]
    val_img = img_data[indices]
    return val_vis, val_img