Update app.py

app.py

@@ -2,56 +2,108 @@ import tensorflow as tf
 import gradio as gr
 import numpy as np
 import cv2
+from scipy.ndimage import center_of_mass
+import math
 
-# 1. Load Model (Optimized load)
-model = tf.keras.models.load_model('TP_MNIST_CNN_model.h5')
-
-def predict_digit(input_data):
-    """
-    Pipeline de prédiction robuste :
-    1. Gestion du format d'entrée (Gradio 4 renvoie parfois un dict).
-    2. Resize vers 28x28 (Interpolation AREA pour préserver les traits).
-    3. Conversion Grayscale.
-    4. Inversion des couleurs (Adaptation domaine Humain -> Machine).
-    5. Normalisation et inférence.
-    """
+# 1. Chargement du modèle
+model = tf.keras.models.load_model('mnist_cnn_v1.keras')
+
+def get_best_shift(img):
+    """Calcule le décalage optimal pour centrer l'image par centre de masse."""
+    cy, cx = center_of_mass(img)
+    rows, cols = img.shape
+    shiftx = np.round(cols/2.0-cx).astype(int)
+    shifty = np.round(rows/2.0-cy).astype(int)
+    return shiftx, shifty
+
+def shift(img, sx, sy):
+    """Applique le décalage géométrique."""
+    rows, cols = img.shape
+    M = np.float32([[1, 0, sx], [0, 1, sy]])
+    shifted = cv2.warpAffine(img, M, (cols, rows))
+    return shifted
+
+def preprocess_image(input_data):
+    """Pipeline robuste : Resize -> Gray -> Invert -> Center -> Normalize"""
     if input_data is None:
         return None
 
-    # Gradio 4 handle : input_data peut être un dictionnaire {'composite': ...}
-    image = input_data["composite"] if isinstance(input_data, dict) else input_data
+    # Gestion format Gradio 4 (Dict ou Array)
+    img = input_data["composite"] if isinstance(input_data, dict) else input_data
+
+    # 1. Resize initial + Grayscale
+    img = cv2.resize(img, (28, 28), interpolation=cv2.INTER_AREA)
+    if len(img.shape) == 3:
+        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
 
-    # Pipeline OpenCV
-    # Resize vers 28x28
-    image = cv2.resize(image, (28, 28), interpolation=cv2.INTER_AREA)
+    # 2. Inversion (Noir sur Blanc -> Blanc sur Noir)
+    # MNIST est blanc sur noir. Si l'utilisateur dessine en noir, on inverse.
+    if np.mean(img) > 127:
+        img = 255 - img
+
+    # 3. Nettoyage du bruit (Thresholding)
+    (_, img) = cv2.threshold(img, 128, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
+
+    # 4. Centrage par Centre de Masse (CRITIQUE pour la précision)
+    # On ajoute une marge noire temporaire pour éviter de couper le chiffre lors du shift
+    while np.sum(img[0]) == 0: img = img[1:]
+    while np.sum(img[:,0]) == 0: img = img[:,1:]
+    while np.sum(img[-1]) == 0: img = img[:-1]
+    while np.sum(img[:,-1]) == 0: img = img[:,:-1]
     
-    # Convertir en niveaux de gris si nécessaire
-    if len(image.shape) == 3:
-        image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+    rows, cols = img.shape
+    if rows > cols:
+        factor = 20.0/rows
+        rows = 20
+        cols = int(round(cols*factor))
+        img = cv2.resize(img, (cols, rows))
+    else:
+        factor = 20.0/cols
+        cols = 20
+        rows = int(round(rows*factor))
+        img = cv2.resize(img, (cols, rows))
+        
+    colsPadding = (int(math.ceil((28-cols)/2.0)), int(math.floor((28-cols)/2.0)))
+    rowsPadding = (int(math.ceil((28-rows)/2.0)), int(math.floor((28-rows)/2.0)))
+    img = np.lib.pad(img, (rowsPadding, colsPadding), 'constant')
+
+    shiftx, shifty = get_best_shift(img)
+    shifted = shift(img, shiftx, shifty)
+    img = shifted
 
-    # Inversion intelligente : Si l'image est majoritairement blanche (dessin noir sur fond blanc)
-    # on inverse car MNIST a été entraîné sur blanc sur fond noir.
-    if np.mean(image) > 127:
-        image = 255 - image
+    # 5. Normalisation et Reshape final
+    img = img / 255.0
+    img = img.reshape(1, 28, 28, 1)
+    return img
 
-    # Normalisation
-    image = image / 255.0
+def predict(image):
+    if image is None: return None
     
-    # Reshape (Batch, H, W, Channels)
-    image = image.reshape(1, 28, 28, 1)
-
-    # Inférence
-    prediction = model.predict(image, verbose=0)
-    return int(np.argmax(prediction))
-
-# Interface Gradio 4 Moderne
-iface = gr.Interface(
-    fn=predict_digit, 
-    inputs=gr.Sketchpad(label="Dessinez un chiffre", type="numpy", crop_size=(28, 28)),
-    outputs="label",
-    title="Reconnaissance MNIST - Production Grade",
-    description="CNN Model. Dessinez un chiffre au centre.",
-    allow_flagging="never"
-)
-
-iface.launch()
+    processed_img = preprocess_image(image)
+    prediction = model.predict(processed_img, verbose=0)[0]
+    
+    # Retourne un dictionnaire {Label: Confiance} pour Gradio
+    return {str(i): float(prediction[i]) for i in range(10)}
+
+# --- UI Moderne avec Gradio Blocks ---
+css = """
+.container {max-width: 800px; margin: auto; padding-top: 20px}
+#title {text-align: center; font-size: 2em; font-weight: bold; margin-bottom: 20px}
+"""
+
+with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
+    gr.Markdown("# 🧠 MNIST Classifier: From Sketch to Prediction", elem_id="title")
+    gr.Markdown("Dessinez un chiffre (0-9). L'IA le centrera automatiquement avant l'analyse.")
+    
+    with gr.Row():
+        with gr.Column():
+            input_sketch = gr.Sketchpad(label="Dessinez ici", type="numpy", crop_size=(200, 200))
+            predict_btn = gr.Button("Analyser", variant="primary")
+        
+        with gr.Column():
+            label_output = gr.Label(num_top_classes=3, label="Prédictions & Probabilités")
+    
+    predict_btn.click(fn=predict, inputs=input_sketch, outputs=label_output)
+
+if __name__ == "__main__":
+    demo.launch()