Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import chess
 import chess.svg
 import io
+import json
 import numpy as np
 import cv2
 import os
@@ -10,11 +11,22 @@ from pathlib import Path
 
 # Import EXACT SAME functions from main.py
 from preprocess import preprocess_image
-from train_tensorflow import create_model
-
-PIECES = ['Empty', 'Rook_White', 'Rook_Black', 'Knight_White', 'Knight_Black', 'Bishop_White',
-          'Bishop_Black', 'Queen_White', 'Queen_Black', 'King_White', 'King_Black', 'Pawn_White', 'Pawn_Black']
-PIECES.sort()
+from train import create_model
+
+# Charger l'ordre des classes depuis le fichier généré par train.py
+try:
+    with open('./class_indices.json', 'r') as f:
+        class_indices = json.load(f)
+    # Inverser pour avoir index -> nom
+    PIECES = [None] * len(class_indices)
+    for name, idx in class_indices.items():
+        PIECES[idx] = name
+    print(f"Ordre des classes chargé: {PIECES}")
+except FileNotFoundError:
+    # Fallback sur ordre alphabétique si le fichier n'existe pas
+    PIECES = ['Bishop_Black', 'Bishop_White', 'Empty', 'King_Black', 'King_White', 'Knight_Black',
+              'Knight_White', 'Pawn_Black', 'Pawn_White', 'Queen_Black', 'Queen_White', 'Rook_Black', 'Rook_White']
+    print(f"Fichier class_indices.json non trouvé, utilisation ordre par défaut")
 
 LABELS = {
     'Empty': '.',
@@ -40,6 +52,14 @@ print("Model loaded!")
 
 
 def classify_image(img):
+    '''Given an image of a single piece, classifies it into one of the classes
+    defined in PIECES.'''
+    # IMPORTANT: Normaliser l'image comme dans l'entraînement (rescale=1/255)
+    if img.max() > 1.0:
+        img = img.astype(np.float32) / 255.0
+    else:
+        img = img.astype(np.float32)
+    
     y_prob = model.predict(img.reshape(1, 300, 150, 3), verbose=0)
     y_pred = y_prob.argmax()
     return PIECES[y_pred]