60 lines
1.9 KiB
Python
60 lines
1.9 KiB
Python
import cv2
|
|
import os
|
|
import numpy as np
|
|
import pickle
|
|
from pathlib import Path
|
|
|
|
# Définir les chemins
|
|
IMAGE_DIR = Path(r"W:/CARIA/images/avatars")
|
|
LABELS_FILE = Path("server-ia/data/modeles/CV2/labels.pickle")
|
|
TRAINER_FILE = Path("server-ia/data/modeles/CV2/trainner.yml")
|
|
|
|
# Initialiser les variables
|
|
current_id = 0
|
|
label_ids = {}
|
|
x_train = []
|
|
y_labels = []
|
|
|
|
def process_images(image_dir):
|
|
global current_id, label_ids, x_train, y_labels
|
|
print("Début du traitement des images...")
|
|
for root, _, files in os.walk(image_dir):
|
|
if files:
|
|
label = Path(root).name
|
|
for file in files:
|
|
if file.lower().endswith(".jpg"):
|
|
path = Path(root) / file
|
|
if label not in label_ids:
|
|
label_ids[label] = current_id
|
|
current_id += 1
|
|
id_ = label_ids[label]
|
|
image = cv2.imread(str(path), cv2.IMREAD_GRAYSCALE)
|
|
if image is not None:
|
|
x_train.append(image)
|
|
y_labels.append(id_)
|
|
print(f"Traitement terminé. {len(x_train)} images chargées.")
|
|
|
|
def save_data():
|
|
print("Sauvegarde des étiquettes...")
|
|
with open(LABELS_FILE, "wb") as f:
|
|
pickle.dump(label_ids, f)
|
|
print(f"Étiquettes sauvegardées dans {LABELS_FILE}.")
|
|
|
|
def train_recognizer():
|
|
print("Entraînement du reconnaisseur...")
|
|
x_train_np = np.array(x_train)
|
|
y_labels_np = np.array(y_labels)
|
|
recognizer = cv2.face.LBPHFaceRecognizer_create()
|
|
recognizer.train(x_train_np, y_labels_np)
|
|
recognizer.save(str(TRAINER_FILE))
|
|
print(f"Modèle entraîné et sauvegardé dans {TRAINER_FILE}.")
|
|
|
|
def main():
|
|
process_images(IMAGE_DIR)
|
|
save_data()
|
|
train_recognizer()
|
|
print("Traitement complet. Les modèles ont été sauvegardés.")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|