77 lines
3.2 KiB
Python
77 lines
3.2 KiB
Python
import tensorflow as tf
|
|
import numpy as np
|
|
import cv2
|
|
import os
|
|
import pandas as pd
|
|
|
|
# Configuration
|
|
size = 60 # Taille des images à laquelle le modèle a été entraîné
|
|
model_path = "server-ia/data/modeles/tensorflow/tf_modele_speed_panneau.keras"
|
|
labels_csv_path = "server-ia/data/modeles/tensorflow/train_labels.csv"
|
|
confidence_threshold = 0.6 # Seuil de confiance pour considérer une prédiction comme valide
|
|
|
|
# Fonction pour charger le mapping des labels depuis le fichier CSV
|
|
def load_labels_mapping(labels_csv_path):
|
|
df = pd.read_csv(labels_csv_path)
|
|
unique_labels = df['label'].unique()
|
|
labels_mapping = {label: idx for idx, label in enumerate(unique_labels)}
|
|
print(f"Mapping des labels chargé : {labels_mapping}")
|
|
return labels_mapping
|
|
|
|
# Chargement du mapping des labels
|
|
labels_mapping = load_labels_mapping(labels_csv_path)
|
|
|
|
# Inversion du mapping pour obtenir l'étiquette à partir de l'index
|
|
reverse_labels_mapping = {v: k for k, v in labels_mapping.items()}
|
|
|
|
# Chargement du modèle
|
|
model = tf.keras.models.load_model(model_path)
|
|
print(f"Modèle chargé depuis {model_path}")
|
|
|
|
# Fonction pour prédire la catégorie d'une image, encadrer la zone reconnue et afficher les résultats
|
|
def predict_and_draw_bounding_box(image_path, model, labels_mapping, size, confidence_threshold):
|
|
# Chargement de l'image originale
|
|
image = cv2.imread(image_path)
|
|
if image is None:
|
|
raise ValueError(f"Erreur de chargement de l'image : {image_path}")
|
|
|
|
# Redimensionnement et normalisation pour la prédiction
|
|
image_resized = cv2.resize(image, (size, size), cv2.INTER_LANCZOS4)
|
|
image_normalized = image_resized.astype('float32') / 255.0 # Normalisation
|
|
image_batch = np.expand_dims(image_normalized, axis=0) # Ajout d'une dimension pour le batch
|
|
|
|
# Prédiction
|
|
predictions = model.predict(image_batch)
|
|
predicted_index = np.argmax(predictions, axis=1)[0]
|
|
confidence = np.max(predictions)
|
|
|
|
# Récupération du label prédictif ou "Inconnu" si la confiance est faible
|
|
if confidence < confidence_threshold:
|
|
predicted_label = "Inconnu"
|
|
print(f"Confiance faible ({confidence:.2f}). La prédiction n'est pas fiable.")
|
|
else:
|
|
predicted_label = reverse_labels_mapping[predicted_index]
|
|
|
|
# Affichage de la prédiction dans la console
|
|
print(f"Image: {image_path}, Prédiction: {predicted_label} (Confiance: {confidence:.2f})")
|
|
|
|
# Encadrement de la zone reconnue (l'ensemble de l'image redimensionnée)
|
|
height, width, _ = image.shape
|
|
cv2.rectangle(image, (0, 0), (width, height), (0, 255, 0), 2) # Couleur verte et épaisseur de 2
|
|
|
|
# Ajout du texte de la prédiction sur l'image
|
|
cv2.putText(image, f'Predicted: {predicted_label}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
|
|
|
|
# Affichage de l'image avec le cadre
|
|
cv2.imshow("Image avec zone reconnue", image)
|
|
cv2.waitKey(0)
|
|
cv2.destroyAllWindows()
|
|
|
|
# Exemple d'utilisation
|
|
image_to_classify = "server-ia/data/images/panneaux/France_road_sign_B14_(10).svg.png" # Remplacez par le chemin exact de l'image à classifier
|
|
|
|
try:
|
|
predict_and_draw_bounding_box(image_to_classify, model, labels_mapping, size, confidence_threshold)
|
|
except ValueError as e:
|
|
print(e)
|