2024102001

from keras import Input
from keras.layers import Conv2D
from keras.layers import MaxPooling2D
from keras.layers import Dropout
from keras.models import Model
from keras.layers import concatenate
from keras.optimizers import Adam
from keras.layers import Conv2DTranspose
from keras.layers import (Conv2D, MaxPooling2D, Dropout, concatenate,
                          Conv2DTranspose, BatchNormalization, Activation)
import os
import numpy as np
import tensorflow as tf
from keras import Input
from keras.layers import Conv2D, MaxPooling2D, Dropout, concatenate, Conv2DTranspose
from keras.models import Model
from keras.optimizers import Adam


def enhancednet(pretrained_weights=None):
    input_shape = (None, None, 1)
    inputs = Input(shape=input_shape, name='input_img')
    conv1 = Conv2D(16, 5, activation='relu', padding='same')(inputs)
    conv1 = BatchNormalization()(conv1)
    conv1 = Activation('relu')(conv1)
    drop1 = Dropout(0.3)(conv1)

    pool1 = MaxPooling2D(pool_size=(2, 2))(drop1)

    conv2 = Conv2D(24, 5, activation='relu', padding='same')(pool1)
    conv2 = BatchNormalization()(conv2)
    conv2 = Activation('relu')(conv2)
    drop2 = Dropout(0.3)(conv2)

    pool2 = MaxPooling2D(pool_size=(2, 2))(drop2)

    conv3 = Conv2D(32, 5, activation='relu', padding='same')(pool2)
    conv3 = BatchNormalization()(conv3)
    conv3 = Activation('relu')(conv3)
    drop3 = Dropout(0.3)(conv3)
    pool3 = MaxPooling2D(pool_size=(2, 2))(drop3)

    conv4 = Conv2D(40, 5, activation='relu', padding='same')(pool3)
    conv4 = BatchNormalization()(conv4)
    conv4 = Activation('relu')(conv4)
    drop4 = Dropout(0.3)(conv4)

    up5 = Conv2D(32, 3, activation='relu', padding='same')(
        Conv2DTranspose(32, 5, activation='relu', padding="same", strides=2)(drop4))
    merge5 = concatenate([drop3, up5], axis=3)
    conv5 = Conv2D(32, 5, activation='relu', padding='same')(merge5)
    drop5 = Dropout(0.6)(conv5)

    up6 = Conv2D(24, 3, activation='relu', padding='same')(
        Conv2DTranspose(24, 5, activation='relu', padding="same", strides=2)(drop5))
    merge6 = concatenate([drop2, up6], axis=3)
    conv6 = Conv2D(24, 5, activation='relu', padding='same')(merge6)
    drop6 = Dropout(0.6)(conv6)

    up7 = Conv2D(16, 3, activation='relu', padding='same')(
        Conv2DTranspose(16, 5, activation='relu', padding="same", strides=2)(drop6))
    merge7 = concatenate([drop1, up7], axis=3)
    conv7 = Conv2D(16, 5, activation='relu', padding='same')(merge7)
    drop7 = Dropout(0.6)(conv7)
    conv8 = Conv2D(1, 1, activation='relu')(drop7)

    model = Model(inputs=inputs, outputs=conv8)

    opt = Adam()

    model.compile(optimizer=opt, loss='mse', metrics=['accuracy'])

    if pretrained_weights:
        model.load_weights(pretrained_weights)

    return model

from dataset import generate_data
import numpy as np

import os
import tensorflow as tf
from model import enhancednet

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

print("可用的GPU设备:")
for gpu in tf.config.list_physical_devices('GPU'):
    print(f"  - {gpu}")


image_rows = 128
image_cols = 256
filename = 'detached_data.mat'
train_data, train1_data, label_data = generate_data(filename)
train_data = np.array(train_data, dtype=float)
train_data = np.reshape(train_data, (len(train_data), image_rows, image_cols, 1))
# print(train_data.shape)
train1_data = np.array(train1_data, dtype=float)
train1_data = np.reshape(train1_data, (len(train1_data), image_rows, image_cols, 1))

model = enhancednet()
model.fit(train_data, train1_data, batch_size=32, epochs=100, verbose=2, shuffle=True, validation_split=0.1)

model.save('enhanced_model.h5')