python
# Импорт библиотек (оставил теми же — это не влияет на результат)
import numpy as np
import matplotlib.pyplot as plt
import segyio
import matplotlib.cm as cm
from keras import layers
import keras


# ---------- Модель (та же архитектура, но собрана через списки блоков) ----------
def get_model(img_size, num_classes):
    init = keras.initializers.HeNormal(seed=14560)
    inputs = keras.Input(shape=img_size + (1,))

    def conv_bn_relu(x, filters):
        x = layers.Conv2D(filters, 3, strides=1, padding="same",
                          kernel_initializer=init)(x)
        x = layers.BatchNormalization()(x)
        return layers.Activation("relu")(x)

    # Encoder: (Conv+BN+ReLU) -> MaxPool (3 раза)
    x = inputs
    for f in (32, 64, 128):
        x = conv_bn_relu(x, f)
        x = layers.MaxPooling2D(2, strides=2, padding="same")(x)

    # Decoder: UpSampling -> (Conv+BN+ReLU) (3 раза)
    for f in (128, 64, 32):
        x = layers.UpSampling2D(2)(x)
        x = conv_bn_relu(x, f)

    outputs = layers.Conv2D(num_classes, 3, activation="softmax",
                            padding="same", kernel_initializer=init)(x)

    return keras.Model(inputs, outputs)


# ---------- Подготовка данных (та же логика, но вынесена в функции) ----------
def _open_segy_pair(seis_path, mask_path):
    seg = segyio.open(seis_path, ignore_geometry=True)
    seg_mask = segyio.open(mask_path, ignore_geometry=True)
    return seg, seg_mask


def _get_inline_indices(seg, train_inline=1000, test_inline=800):
    inline = seg.attributes(segyio.TraceField.INLINE_3D)[:]
    ind_train = np.where(inline == train_inline)[0]
    ind_test = np.where(inline == test_inline)[0]
    return ind_train, ind_test


def _alloc_arrays(h=1400, w=160):
    S_train = np.zeros((1, h, w, 1))
    S_test = np.zeros((1, h, w, 1))
    M_train = np.zeros((1, h, w, 2))
    M_test = np.zeros((1, h, w, 2))
    return S_train, S_test, M_train, M_test


def _fill_arrays(seg, seg_mask, ind_train, ind_test, S_train, S_test, M_train, M_test,
                 h=1400, w=160, offset=400):
    # Эквивалентно исходнику: те же циклы, те же индексы, те же присваивания
    for i in range(h):
        tr_tr = ind_train[i]
        tr_te = ind_test[i]
        for k in range(w):
            idx = offset + k

            s_tr = seg.trace[tr_tr][idx]
            s_te = seg.trace[tr_te][idx]

            m_tr = seg_mask.trace[tr_tr][idx]
            m_te = seg_mask.trace[tr_te][idx]

            S_train[0, i, k, 0] = s_tr
            S_test[0, i, k, 0] = s_te

            M_train[0, i, k, 0] = m_tr
            M_train[0, i, k, 1] = 1 - m_tr

            M_test[0, i, k, 0] = m_te
            M_test[0, i, k, 1] = 1 - m_te


# ---------- Визуализация (та же, но через словарь) ----------
def _plot_results(S_train, M_train, Mask_train, S_test, M_test, Mask_test):
    fig, ax = plt.subplots(2, 3, figsize=(14, 7))

    panels = [
        (0, 0, S_train[0, :, :, 0].T, cm.seismic, "Train: Seismic"),
        (0, 1, M_train[0, :, :, 0].T, cm.binary,  "Train: True mask"),
        (0, 2, Mask_train[0, :, :, 0].T, cm.binary, "Train: Pred mask"),
        (1, 0, S_test[0, :, :, 0].T, cm.seismic, "Test: Seismic"),
        (1, 1, M_test[0, :, :, 0].T, cm.binary,  "Test: True mask"),
        (1, 2, Mask_test[0, :, :, 0].T, cm.binary, "Test: Pred mask"),
    ]

    for r, c, data, cmap, title in panels:
        ax[r, c].imshow(data, cmap=cmap, aspect="auto")
        ax[r, c].set_title(title)

    for a in ax.ravel():
        a.set_xlabel("Trace index")
        a.set_ylabel("Sample (window)")

    plt.show()


# ---------- Основной запуск (та же последовательность действий) ----------
def run_model_and_visualize():
    try:
        seg, seg_mask = _open_segy_pair("Cube.sgy", "Biogerm_mask.sgy")
        ind_train, ind_test = _get_inline_indices(seg, train_inline=1000, test_inline=800)

        S_train, S_test, M_train, M_test = _alloc_arrays(h=1400, w=160)
        _fill_arrays(seg, seg_mask, ind_train, ind_test,
                     S_train, S_test, M_train, M_test,
                     h=1400, w=160, offset=400)

        model = get_model((1400, 160), 2)
        model.compile(optimizer="rmsprop", loss="binary_crossentropy")
        model.fit(S_train, M_train, epochs=150, verbose=0)

        Mask_train = model.predict(S_train, verbose=0)
        Mask_test = model.predict(S_test, verbose=0)

        _plot_results(S_train, M_train, Mask_train, S_test, M_test, Mask_test)

    except Exception as e:
        print(f"Error: {e}")


# Запуск модели и визуализация
run_model_and_visualize()
# Код запускается примерно 20 минут с визуализацией результатов