You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
106 lines
4.2 KiB
106 lines
4.2 KiB
import os
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D, Concatenate, Dropout
|
|
from tensorflow.keras.models import Model
|
|
from tensorflow.keras.optimizers import Adam
|
|
from tensorflow.keras.preprocessing.image import img_to_array, load_img
|
|
from sklearn.model_selection import train_test_split
|
|
|
|
# Charger les données UTKFace
|
|
def load_utkface_dataset(dataset_path):
|
|
images = []
|
|
ages = []
|
|
for file_name in os.listdir(dataset_path):
|
|
if file_name.endswith(".jpg"):
|
|
age = int(file_name.split("_")[0]) # Extraire l'âge du nom du fichier
|
|
img_path = os.path.join(dataset_path, file_name)
|
|
img = load_img(img_path, target_size=(128, 128))
|
|
img_array = img_to_array(img) / 255.0
|
|
images.append(img_array)
|
|
ages.append(age)
|
|
return np.array(images), np.array(ages)
|
|
|
|
# Définir le modèle U-Net
|
|
def build_unet(input_shape):
|
|
inputs = Input(shape=input_shape)
|
|
|
|
# Encoder
|
|
c1 = Conv2D(64, (3, 3), activation='relu', padding='same')(inputs)
|
|
c1 = Dropout(0.1)(c1)
|
|
c1 = Conv2D(64, (3, 3), activation='relu', padding='same')(c1)
|
|
p1 = MaxPooling2D((2, 2))(c1)
|
|
|
|
c2 = Conv2D(128, (3, 3), activation='relu', padding='same')(p1)
|
|
c2 = Dropout(0.1)(c2)
|
|
c2 = Conv2D(128, (3, 3), activation='relu', padding='same')(c2)
|
|
p2 = MaxPooling2D((2, 2))(c2)
|
|
|
|
c3 = Conv2D(256, (3, 3), activation='relu', padding='same')(p2)
|
|
c3 = Dropout(0.2)(c3)
|
|
c3 = Conv2D(256, (3, 3), activation='relu', padding='same')(c3)
|
|
p3 = MaxPooling2D((2, 2))(c3)
|
|
|
|
c4 = Conv2D(512, (3, 3), activation='relu', padding='same')(p3)
|
|
c4 = Dropout(0.2)(c4)
|
|
c4 = Conv2D(512, (3, 3), activation='relu', padding='same')(c4)
|
|
p4 = MaxPooling2D((2, 2))(c4)
|
|
|
|
c5 = Conv2D(1024, (3, 3), activation='relu', padding='same')(p4)
|
|
c5 = Dropout(0.3)(c5)
|
|
c5 = Conv2D(1024, (3, 3), activation='relu', padding='same')(c5)
|
|
|
|
# Decoder
|
|
u6 = UpSampling2D((2, 2))(c5)
|
|
u6 = Conv2D(512, (2, 2), activation='relu', padding='same')(u6)
|
|
u6 = Concatenate()([u6, c4])
|
|
c6 = Conv2D(512, (3, 3), activation='relu', padding='same')(u6)
|
|
c6 = Dropout(0.2)(c6)
|
|
c6 = Conv2D(512, (3, 3), activation='relu', padding='same')(c6)
|
|
|
|
u7 = UpSampling2D((2, 2))(c6)
|
|
u7 = Conv2D(256, (2, 2), activation='relu', padding='same')(u7)
|
|
u7 = Concatenate()([u7, c3])
|
|
c7 = Conv2D(256, (3, 3), activation='relu', padding='same')(u7)
|
|
c7 = Dropout(0.2)(c7)
|
|
c7 = Conv2D(256, (3, 3), activation='relu', padding='same')(c7)
|
|
|
|
u8 = UpSampling2D((2, 2))(c7)
|
|
u8 = Conv2D(128, (2, 2), activation='relu', padding='same')(u8)
|
|
u8 = Concatenate()([u8, c2])
|
|
c8 = Conv2D(128, (3, 3), activation='relu', padding='same')(u8)
|
|
c8 = Dropout(0.1)(c8)
|
|
c8 = Conv2D(128, (3, 3), activation='relu', padding='same')(c8)
|
|
|
|
u9 = UpSampling2D((2, 2))(c8)
|
|
u9 = Conv2D(64, (2, 2), activation='relu', padding='same')(u9)
|
|
u9 = Concatenate()([u9, c1])
|
|
c9 = Conv2D(64, (3, 3), activation='relu', padding='same')(u9)
|
|
c9 = Dropout(0.1)(c9)
|
|
c9 = Conv2D(64, (3, 3), activation='relu', padding='same')(c9)
|
|
|
|
outputs = Conv2D(3, (1, 1), activation='sigmoid')(c9)
|
|
|
|
model = Model(inputs, outputs)
|
|
return model
|
|
|
|
# Entraîner le modèle U-Net
|
|
def train_unet(model, X_train, X_test, epochs, batch_size):
|
|
model.compile(optimizer=Adam(learning_rate=1e-4), loss='mean_squared_error', metrics=['accuracy'])
|
|
model.fit(X_train, X_train, validation_data=(X_test, X_test), epochs=epochs, batch_size=batch_size)
|
|
model.save("unet_face_aging_model.h5")
|
|
print("Modèle entraîné et sauvegardé avec succès !")
|
|
|
|
if __name__ == "__main__":
|
|
dataset_path = "./UTKFace/part1/part1" # Modifier selon l'emplacement du dataset
|
|
images, ages = load_utkface_dataset(dataset_path)
|
|
|
|
# Diviser les données en ensembles d'entraînement et de test
|
|
X_train, X_test, _, _ = train_test_split(images, ages, test_size=0.2, random_state=42)
|
|
|
|
# Construire le modèle U-Net
|
|
input_shape = (128, 128, 3)
|
|
unet_model = build_unet(input_shape)
|
|
|
|
# Entraîner le modèle U-Net
|
|
train_unet(unet_model, X_train, X_test, epochs=100, batch_size=32) |