221 lines
8.9 KiB
Python
221 lines
8.9 KiB
Python
import tensorflow as tf
|
|
from tensorflow.keras import layers, models
|
|
import numpy as np
|
|
from sklearn.model_selection import train_test_split
|
|
import cv2
|
|
import os
|
|
import time
|
|
from tqdm import tqdm
|
|
|
|
# Fonction pour charger les labels depuis un fichier CSV
|
|
def load_labels_from_csv(filepath):
|
|
filenames, labels = [], []
|
|
with open(filepath, 'r') as file:
|
|
next(file) # Skip header
|
|
for line in file:
|
|
parts = line.strip().split(',')
|
|
if len(parts) == 2:
|
|
filenames.append(parts[0])
|
|
labels.append(parts[1])
|
|
return filenames, labels
|
|
|
|
# Chargement des fichiers CSV des labels
|
|
panneaux_filenames, panneaux_labels = load_labels_from_csv('server-trainer/images/vitesse_panneaux_labels.csv')
|
|
autres_filenames, autres_labels = load_labels_from_csv('server-trainer/images/autres_panneaux_labels.csv')
|
|
sans_filenames, sans_labels = load_labels_from_csv('server-trainer/images/genere_sans_panneaux_labels.csv')
|
|
|
|
# Création d'un dictionnaire pour les labels
|
|
def create_label_dict(filenames, labels):
|
|
label_dict = {i: label for i, label in enumerate(labels)}
|
|
return label_dict
|
|
|
|
# Dictionnaires de labels
|
|
panneaux_labels_dict = create_label_dict(panneaux_filenames, panneaux_labels)
|
|
autres_labels_dict = create_label_dict(autres_filenames, autres_labels)
|
|
sans_labels_dict = {0: "Sans panneau"}
|
|
|
|
# Configuration
|
|
size = 60
|
|
dir_images_panneaux = "server-trainer/images/vitesse_panneaux"
|
|
dir_images_autres_panneaux = "server-trainer/images/autres_panneaux"
|
|
dir_images_genere_sans_panneaux = "server-trainer/images/genere_sans_panneaux"
|
|
batch_size = 128
|
|
nbr_entrainement = 1 # Nombre d'époques d'entraînement
|
|
|
|
print(f"Configuration : taille des images = {size}, taille du batch = {batch_size}, nombre d'époques = {nbr_entrainement}")
|
|
|
|
# Définition du modèle
|
|
def panneau_model(nbr_classes):
|
|
print(f"Création du modèle avec {nbr_classes} classes.")
|
|
model = tf.keras.Sequential([
|
|
layers.Input(shape=(size, size, 3)),
|
|
layers.Conv2D(128, 3, padding='same', activation='relu'),
|
|
layers.Dropout(0.2),
|
|
layers.BatchNormalization(),
|
|
layers.Conv2D(128, 3, padding='same', activation='relu'),
|
|
layers.Dropout(0.2),
|
|
layers.BatchNormalization(),
|
|
layers.MaxPool2D(pool_size=2),
|
|
layers.Conv2D(256, 3, padding='same', activation='relu'),
|
|
layers.Dropout(0.3),
|
|
layers.BatchNormalization(),
|
|
layers.Conv2D(256, 3, padding='same', activation='relu'),
|
|
layers.Dropout(0.4),
|
|
layers.BatchNormalization(),
|
|
layers.MaxPool2D(pool_size=2),
|
|
layers.Flatten(),
|
|
layers.Dense(512, activation='relu'),
|
|
layers.Dropout(0.5),
|
|
layers.BatchNormalization(),
|
|
layers.Dense(nbr_classes, activation='sigmoid')
|
|
])
|
|
print("Modèle créé.")
|
|
return model
|
|
|
|
# Chargement des images depuis un répertoire
|
|
def load_images_from_directory(directory, size):
|
|
if not os.path.exists(directory):
|
|
raise FileNotFoundError(f"Le répertoire d'images n'existe pas : {directory}")
|
|
|
|
files = [f for f in sorted(os.listdir(directory)) if f.endswith(".png")]
|
|
if not files:
|
|
raise FileNotFoundError(f"Le répertoire d'images est vide : {directory}")
|
|
|
|
images = []
|
|
for file in files:
|
|
path = os.path.join(directory, file)
|
|
image = cv2.imread(path)
|
|
image = cv2.resize(image, (size, size), cv2.INTER_LANCZOS4)
|
|
images.append(image)
|
|
|
|
return images
|
|
|
|
# Préparation des données
|
|
def prepare_data(panneaux_images, autres_images, sans_panneaux_images, panneau_labels_dict):
|
|
print("Préparation des données pour les images de panneaux...")
|
|
tab_images = np.array([])
|
|
tab_labels = np.array([])
|
|
|
|
# Images de panneaux
|
|
for i, image in enumerate(panneaux_images):
|
|
lot = np.array([image for _ in range(120)])
|
|
if tab_images.size == 0:
|
|
tab_images = lot
|
|
else:
|
|
tab_images = np.concatenate((tab_images, lot), axis=0)
|
|
|
|
labels = np.eye(len(panneaux_labels_dict))[i]
|
|
labels = np.repeat([labels], len(lot), axis=0)
|
|
if tab_labels.size == 0:
|
|
tab_labels = labels
|
|
else:
|
|
tab_labels = np.concatenate([tab_labels, labels], axis=0)
|
|
|
|
print(f"Nombre d'images après ajout des panneaux : {len(tab_images)}")
|
|
print(f"Nombre de labels après ajout des panneaux : {len(tab_labels)}")
|
|
|
|
# Images des autres panneaux
|
|
print("Traitement des autres panneaux...")
|
|
for image in autres_images:
|
|
lot = np.array([image for _ in range(700)])
|
|
if tab_images.size == 0:
|
|
tab_images = lot
|
|
else:
|
|
tab_images = np.concatenate([tab_images, lot], axis=0)
|
|
|
|
labels = np.zeros((len(lot), len(panneaux_labels_dict)))
|
|
if tab_labels.size == 0:
|
|
tab_labels = labels
|
|
else:
|
|
tab_labels = np.concatenate([tab_labels, labels], axis=0)
|
|
|
|
print(f"Nombre d'images après ajout des autres panneaux : {len(tab_images)}")
|
|
print(f"Nombre de labels après ajout des autres panneaux : {len(tab_labels)}")
|
|
|
|
# Images générées sans panneaux
|
|
print("Traitement des images générées sans panneaux...")
|
|
sans_panneaux_images = [cv2.resize(img, (size, size)) for img in sans_panneaux_images]
|
|
tab_images = np.concatenate([tab_images, np.array(sans_panneaux_images)], axis=0)
|
|
tab_labels = np.concatenate([tab_labels, np.zeros((len(sans_panneaux_images), len(panneaux_labels_dict)))], axis=0)
|
|
|
|
print(f"Nombre d'images après ajout des images sans panneaux : {len(tab_images)}")
|
|
print(f"Nombre de labels après ajout des images sans panneaux : {len(tab_labels)}")
|
|
|
|
# Normalisation
|
|
tab_images = tab_images.astype(np.float32) / 255
|
|
tab_labels = tab_labels.astype(np.float32)
|
|
|
|
return tab_images, tab_labels
|
|
|
|
# Chargement des images
|
|
panneaux_images = load_images_from_directory(dir_images_panneaux, size)
|
|
autres_images = load_images_from_directory(dir_images_autres_panneaux, size)
|
|
sans_panneaux_images = load_images_from_directory(dir_images_genere_sans_panneaux, size)
|
|
|
|
# Préparation des données
|
|
tab_images, tab_labels = prepare_data(panneaux_images, autres_images, sans_panneaux_images, panneaux_labels_dict)
|
|
|
|
# Division du jeu de données
|
|
train_images, test_images, train_labels, test_labels = train_test_split(tab_images, tab_labels, test_size=0.10, random_state=42)
|
|
|
|
print(f"Nombre d'images d'entraînement : {len(train_images)}, Nombre d'images de test : {len(test_images)}")
|
|
|
|
train_ds = tf.data.Dataset.from_tensor_slices((train_images, train_labels)).batch(batch_size)
|
|
test_ds = tf.data.Dataset.from_tensor_slices((test_images, test_labels)).batch(batch_size)
|
|
|
|
# Compilation du modèle
|
|
model_panneau = panneau_model(len(panneaux_labels_dict))
|
|
model_panneau.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
|
|
print("Modèle compilé.")
|
|
|
|
# Entraînement du modèle avec barre de progression pour chaque époque
|
|
def train(model, train_ds, test_ds, nbr_entrainement):
|
|
for epoch in range(nbr_entrainement):
|
|
print(f"Début de l'entraînement {epoch + 1}...")
|
|
start = time.time()
|
|
|
|
# Nombre de batches dans une époque
|
|
num_batches = len(train_ds)
|
|
|
|
# Initialiser tqdm pour les batches
|
|
with tqdm(total=num_batches, desc=f"Époque {epoch + 1}", unit='batch') as pbar:
|
|
for batch_images, batch_labels in train_ds:
|
|
history = model.train_on_batch(batch_images, batch_labels)
|
|
pbar.update(1) # Mise à jour de la barre de progression
|
|
|
|
# Afficher les résultats après l'époque
|
|
train_loss, train_accuracy = history
|
|
print(f'Entraînement {epoch + 1:04d} : perte : {train_loss:6.4f}, précision : {train_accuracy * 100:7.4f}%, temps : {time.time() - start:7.4f}')
|
|
test(model, test_ds)
|
|
|
|
def test(model, test_ds):
|
|
print("Début du test...")
|
|
start = time.time()
|
|
test_loss, test_accuracy = model.evaluate(test_ds, verbose=0)
|
|
print(f' >>> Test : perte : {test_loss:6.4f}, précision : {test_accuracy * 100:7.4f}%, temps : {time.time() - start:7.4f}')
|
|
|
|
print("Début de l'entraînement du modèle.")
|
|
train(model_panneau, train_ds, test_ds, nbr_entrainement)
|
|
|
|
model_panneau.save("server-ia/data/modeles/tensorflow/tf_modele_AllRoadSign.keras")
|
|
print(f"Modèle sauvegardé à : {os.path.abspath(model_panneau)}")
|
|
|
|
# Évaluation des prédictions
|
|
print("Évaluation des prédictions sur les images de test.")
|
|
for i in range(len(test_images)):
|
|
prediction = model_panneau.predict(np.array([test_images[i]]))
|
|
predicted_index = np.argmax(prediction[0])
|
|
confidence = np.max(prediction[0])
|
|
|
|
# Affichage du label correspondant
|
|
if confidence < 0.6:
|
|
print("Ce n'est pas un panneau")
|
|
else:
|
|
print(f"C'est un panneau : {panneaux_labels_dict.get(predicted_index, 'Inconnu')}")
|
|
|
|
# Affichage de l'image
|
|
cv2.imshow("image", test_images[i])
|
|
if cv2.waitKey() & 0xFF == ord('q'):
|
|
break
|
|
cv2.destroyAllWindows()
|