datagenerator.Augmentation

  1from scipy.interpolate import PchipInterpolator
  2from scipy.interpolate import UnivariateSpline
  3import numpy as np
  4import math
  5from scipy.ndimage import uniform_filter
  6from scipy.ndimage.morphology import distance_transform_edt
  7
  8
  9def tz_stretch(x_cube, y_cube, wb):
 10    """Apply stretch/squeeze to simulate time-depth conversion
 11
 12    Arguments:
 13        x_cube {ndarray} -- 4D array containing seismic data, 1st dimension angles
 14        y_cube {ndarray} -- 3D array containing hydrocarbon labels
 15        wb {ndarray} -- 2D array containing seabed horizon in samples
 16
 17    Returns:
 18        [ndarray, ndarray] -- Stretched/squeezed seismic data and hydrocarbon labels
 19    """
 20    n_cubes, x_shape, y_shape, z_shape = x_cube.shape
 21    seafloor_map = uniform_filter(wb, size=(31, 31))
 22    seafloor_min = int(seafloor_map.min())
 23
 24    max_squeeze_depth = np.random.triangular(
 25        seafloor_map.min() * 1.2,
 26        (seafloor_map.min() + z_shape) / 2.0,
 27        z_shape * 0.8,
 28    )
 29    max_squeeze_pct = np.random.triangular(1.05, 1.15, 1.25)
 30
 31    # build function between ascending indices and stretched/squeezed indices
 32    # - hang from min depth for seafloor
 33    raw_depth_indices = np.arange(0.0, z_shape)
 34    stretched_depth_indices = raw_depth_indices.copy()
 35
 36    # get depth (as index) and percentage to squeeze point that is midway
 37    # between seafloor and bottom of trace
 38    fit_indices = np.array(
 39        [seafloor_min, seafloor_min + 1.0, max_squeeze_depth, z_shape - 1, z_shape]
 40    )
 41    fit_depths = np.array(
 42        (
 43            seafloor_min,
 44            seafloor_min + 1.0,
 45            max_squeeze_depth / max_squeeze_pct,
 46            z_shape - 1,
 47            z_shape,
 48        )
 49    )
 50    spline_interp = UnivariateSpline(fit_indices, fit_depths)
 51    stretched_depth_indices[seafloor_min + 1 : z_shape - 1] = spline_interp(
 52        raw_depth_indices[seafloor_min + 1 : z_shape - 1]
 53    )
 54
 55    output_x_cube = np.zeros_like(x_cube)
 56    output_y_cube = np.zeros_like(y_cube)
 57    for k in range(n_cubes):
 58        for i in range(x_shape):
 59            for j in range(y_shape):
 60                stretch_times = raw_depth_indices.copy()
 61                fit_indices_ij = fit_indices - seafloor_min + seafloor_map[i, j]
 62                fit_depths_ij = fit_depths - seafloor_min + seafloor_map[i, j]
 63                seafloor_ind = int(seafloor_map[i, j]) + 1
 64                spline_interp = UnivariateSpline(fit_depths_ij, fit_indices_ij)
 65                stretch_times[seafloor_ind : z_shape - 1] = spline_interp(
 66                    raw_depth_indices[seafloor_ind : z_shape - 1]
 67                )
 68                output_trace = np.interp(
 69                    stretch_times, raw_depth_indices, x_cube[k, i, j, :]
 70                )
 71                output_trace[stretch_times >= z_shape - 1.0] = 0.0
 72                output_x_cube[k, i, j, :] = output_trace.copy()
 73
 74                if k == 0:
 75                    # Apply to labels
 76                    output_trace = np.interp(
 77                        stretch_times, raw_depth_indices, y_cube[i, j, :]
 78                    )
 79                    output_trace[stretch_times >= z_shape - 1.0] = 0.0
 80                    output_y_cube[i, j, :] = output_trace.copy()
 81
 82    return output_x_cube, output_y_cube
 83
 84
 85def uniform_stretch(x_cube, y_cube):
 86    x_max, y_max, t_max = x_cube.shape[1:]
 87
 88    x_in = np.arange(x_max)
 89    y_in = np.arange(y_max)
 90    t_in = np.arange(t_max)
 91
 92    squeeze_min = 0.85
 93    squeeze_max = 1.15
 94    # Uniform r.v. squeeze factors
 95    x_squeeze_factor = np.random.rand() * (squeeze_max - squeeze_min) + squeeze_min
 96    y_squeeze_factor = np.random.rand() * (squeeze_max - squeeze_min) + squeeze_min
 97    t_squeeze_factor = np.random.rand() * (squeeze_max - squeeze_min) + squeeze_min
 98
 99    x_out = x_in * x_squeeze_factor
100    if x_squeeze_factor > 1:
101        roll_x_by = math.floor(np.sum(x_out > x_max - 1) / 2)
102        # print(roll_x_by)
103        x_out = np.roll(x_out, roll_x_by)
104
105    y_out = y_in * y_squeeze_factor
106    if y_squeeze_factor > 1:
107        roll_y_by = math.floor(np.sum(y_out > y_max - 1) / 2)
108        # print(roll_y_by)
109        y_out = np.roll(y_out, roll_y_by)
110
111    t_out = t_in * t_squeeze_factor
112    if t_squeeze_factor > 1:
113        roll_t_by = math.floor(np.sum(t_out > t_max - 1) / 2)
114        # print(roll_t_by)
115        t_out = np.roll(t_out, roll_t_by)
116
117    x_cube_sqz = np.apply_along_axis(
118        lambda w: np.interp(x_out, x_in, w, left=0, right=0), axis=1, arr=x_cube
119    )
120    y_cube_sqz = np.apply_along_axis(
121        lambda w: np.interp(x_out, x_in, w, left=0, right=0), axis=0, arr=y_cube
122    )
123    y_cube_sqz = np.round(y_cube_sqz)
124
125    x_cube_sqz = np.apply_along_axis(
126        lambda w: np.interp(y_out, y_in, w, left=0, right=0), axis=2, arr=x_cube_sqz
127    )
128    y_cube_sqz = np.apply_along_axis(
129        lambda w: np.interp(y_out, y_in, w, left=0, right=0), axis=1, arr=y_cube_sqz
130    )
131    y_cube_sqz = np.round(y_cube_sqz)
132
133    x_cube_sqz = np.apply_along_axis(
134        lambda w: np.interp(t_out, t_in, w, left=0, right=0), axis=3, arr=x_cube_sqz
135    )
136    y_cube_sqz = np.apply_along_axis(
137        lambda w: np.interp(t_out, t_in, w, left=0, right=0), axis=2, arr=y_cube_sqz
138    )
139    y_cube_sqz = np.round(y_cube_sqz).astype(
140        "uint8"
141    )  # convert to uint8 for writing to disk
142
143    return x_cube_sqz, y_cube_sqz