Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,135 +1,176 @@
 import json
+from pathlib import Path
+
 import numpy as np
 import streamlit as st
 from PIL import Image
+
 import tensorflow as tf
-from pathlib import Path
 
-st.set_page_config(page_title='Facial Keypoints Predictor (CNN)', page_icon='🙂', layout='centered')
 
+# -------------------------
+# Page config
+# -------------------------
+st.set_page_config(
+    page_title='Facial Keypoints Predictor (CNN)',
+    page_icon='🙂',
+    layout='centered'
+)
+
+st.title('🙂 Facial Keypoints Predictor (CNN)')
+st.write('Upload a face image and the model will predict 15 facial keypoints (30 values: x/y).')
+
+
+# -------------------------
+# Paths (HuggingFace friendly)
+# Put ALL files inside /src
+# -------------------------
 BASE_DIR = Path(__file__).resolve().parent
 
-# --- Choose ONE of these model paths depending on what you uploaded ---
-MODEL_SAVEDMODEL_DIR = BASE_DIR / 'final_keypoints_cnn_savedmodel'   # folder (recommended)
-MODEL_H5_PATH = BASE_DIR / 'final_keypoints_cnn.h5'                 # file (alternative)
+MODEL_KERAS_PATH = BASE_DIR / 'final_keypoints_cnn.keras'
+MODEL_H5_PATH = BASE_DIR / 'final_keypoints_cnn.h5'
 
 TARGET_COLS_PATH = BASE_DIR / 'target_cols.json'
-PREPROCESS_CFG_PATH = BASE_DIR / 'preprocess_config.json'
+PREPROCESS_PATH = BASE_DIR / 'preprocess_config.json'
 
 
+# -------------------------
+# Load assets
+# -------------------------
 @st.cache_resource
 def load_assets():
-    # load metadata
+    # ✅ IMPORTANT: Keras 3 does NOT load SavedModel folders via load_model()
+    # So we FORCE .keras or .h5 only.
+    if MODEL_KERAS_PATH.exists():
+        model = tf.keras.models.load_model(str(MODEL_KERAS_PATH), compile=False)
+        model_source = MODEL_KERAS_PATH.name
+    elif MODEL_H5_PATH.exists():
+        model = tf.keras.models.load_model(str(MODEL_H5_PATH), compile=False)
+        model_source = MODEL_H5_PATH.name
+    else:
+        raise FileNotFoundError(
+            'Model not found. Upload `final_keypoints_cnn.keras` (recommended) or `final_keypoints_cnn.h5` into /src.'
+        )
+
     if not TARGET_COLS_PATH.exists():
-        raise FileNotFoundError(f'Missing {TARGET_COLS_PATH.name} in repo root.')
+        raise FileNotFoundError('Missing file: target_cols.json (put it in /src)')
+
+    if not PREPROCESS_PATH.exists():
+        raise FileNotFoundError('Missing file: preprocess_config.json (put it in /src)')
 
     with open(TARGET_COLS_PATH, 'r') as f:
         target_cols = json.load(f)
 
-    preprocess = {'img_size': [96, 96], 'normalize': 'x / 255.0', 'target_normalization': '(y - 48) / 48'}
-    if PREPROCESS_CFG_PATH.exists():
-        with open(PREPROCESS_CFG_PATH, 'r') as f:
-            preprocess = json.load(f)
+    with open(PREPROCESS_PATH, 'r') as f:
+        preprocess_cfg = json.load(f)
 
-    # load model (prefer SavedModel folder)
-    if MODEL_SAVEDMODEL_DIR.exists():
-        model = tf.keras.models.load_model(str(MODEL_SAVEDMODEL_DIR), compile=False)
-    elif MODEL_H5_PATH.exists():
-        model = tf.keras.models.load_model(str(MODEL_H5_PATH), compile=False)
-    else:
-        raise FileNotFoundError(
-            'Model not found. Upload either:\n'
-            f'- {MODEL_SAVEDMODEL_DIR.name}/ (SavedModel folder)\n'
-            f'- {MODEL_H5_PATH.name} (H5 file)\n'
-        )
+    return model, target_cols, preprocess_cfg, model_source
 
-    return model, target_cols, preprocess
-
-
-def preprocess_image(pil_img, img_size=(96, 96)):
-    img = pil_img.convert('L').resize(img_size)   # grayscale
-    x = np.array(img, dtype=np.float32)
-    x = x / 255.0
-    x = np.expand_dims(x, axis=-1)  # (H, W, 1)
-    x = np.expand_dims(x, axis=0)   # (1, H, W, 1)
-    return x, np.array(img)
-
-
-def denormalize_and_clip(y_pred):
-    # your training normalization: y_norm = (y - 48) / 48
-    # so inverse: y = y_norm * 48 + 48
-    y = y_pred * 48.0 + 48.0
-    y = np.clip(y, 0.0, 96.0)
-    return y
-
-
-def keypoints_to_xy(y_vec, target_cols):
-    # y_vec: shape (30,)
-    coords = {}
-    for i, name in enumerate(target_cols):
-        coords[name] = float(y_vec[i])
-    return coords
-
-
-def draw_points_on_image(gray_img_96, coords, point_size=2):
-    # gray_img_96: (96,96) uint8
-    rgb = np.stack([gray_img_96]*3, axis=-1).copy()
-    # draw red dots
-    for k, v in coords.items():
-        if k.endswith('_x'):
-            y_name = k.replace('_x', '_y')
-            if y_name in coords:
-                x = int(round(coords[k]))
-                y = int(round(coords[y_name]))
-                if 0 <= x < 96 and 0 <= y < 96:
-                    x0, x1 = max(0, x-point_size), min(95, x+point_size)
-                    y0, y1 = max(0, y-point_size), min(95, y+point_size)
-                    rgb[y0:y1+1, x0:x1+1, 0] = 255
-                    rgb[y0:y1+1, x0:x1+1, 1] = 0
-                    rgb[y0:y1+1, x0:x1+1, 2] = 0
-    return rgb
 
+# -------------------------
+# Helpers
+# -------------------------
+def preprocess_image(pil_img: Image.Image, img_size=(96, 96)) -> np.ndarray:
+    # Convert to grayscale like the Kaggle dataset (96x96, 1 channel)
+    img = pil_img.convert('L').resize(img_size)
+    arr = np.array(img).astype(np.float32) / 255.0  # normalize x / 255
+    arr = np.expand_dims(arr, axis=-1)  # (96, 96, 1)
+    arr = np.expand_dims(arr, axis=0)   # (1, 96, 96, 1)
+    return arr
 
-st.title('🙂 Facial Keypoints Predictor (CNN)')
-st.write('Upload a face image and the model will predict 15 facial keypoints (30 values: x/y).')
 
+def draw_keypoints(pil_img: Image.Image, keypoints_xy: np.ndarray) -> Image.Image:
+    # keypoints_xy shape: (15, 2) -> x,y
+    import PIL.ImageDraw as ImageDraw
+
+    img = pil_img.convert('RGB').resize((96, 96))
+    draw = ImageDraw.Draw(img)
+
+    for (x, y) in keypoints_xy:
+        r = 2
+        draw.ellipse((x - r, y - r, x + r, y + r), outline='red', width=2)
+    return img
+
+
+def to_xy(pred_30: np.ndarray) -> np.ndarray:
+    # pred_30 shape: (30,)
+    pts = pred_30.reshape(-1, 2)
+    return pts
+
+
+# -------------------------
+# UI: checklist
+# -------------------------
 with st.expander('Model files checklist'):
     st.markdown(
-        '- **target_cols.json** ✅\n'
-        '- **preprocess_config.json** ✅ (optional)\n'
-        '- **Model**: `final_keypoints_cnn_savedmodel/` (recommended) OR `final_keypoints_cnn.h5`\n'
+        '- Put files inside **`/src`** in your HuggingFace Space.\n'
+        '- Required:\n'
+        '  - `final_keypoints_cnn.keras` (recommended) **or** `final_keypoints_cnn.h5`\n'
+        '  - `target_cols.json`\n'
+        '  - `preprocess_config.json`\n'
+        '- Optional: `history.pkl` (not needed for inference)\n'
+        '\n'
+        '✅ Tip: If you still have a folder `final_keypoints_cnn_savedmodel/`, remove it or ignore it. '
+        'This app does **not** load SavedModel folders.'
     )
 
-model, target_cols, preprocess_cfg = load_assets()
 
+# -------------------------
+# Load model + configs
+# -------------------------
+try:
+    model, target_cols, preprocess_cfg, model_source = load_assets()
+    st.success(f'Model loaded: {model_source}')
+except Exception as e:
+    st.error(str(e))
+    st.stop()
+
+
+# -------------------------
+# Upload + Predict
+# -------------------------
 uploaded = st.file_uploader('Upload an image (jpg/png)', type=['jpg', 'jpeg', 'png'])
 
 if uploaded is not None:
     pil_img = Image.open(uploaded)
-    st.image(pil_img, caption='Uploaded image', use_container_width=True)
+    st.subheader('Input image')
+    st.image(pil_img, use_container_width=True)
+
+    x = preprocess_image(pil_img, img_size=(96, 96))
 
-    x, gray_96 = preprocess_image(pil_img, img_size=tuple(preprocess_cfg.get('img_size', [96, 96])))
+    # Predict
+    pred = model.predict(x, verbose=0)[0]  # shape (30,)
 
-    y_pred_norm = model.predict(x, verbose=0)[0]          # (30,)
-    y_pred = denormalize_and_clip(y_pred_norm)           # (30,)
+    # If your model predicts normalized coordinates, you must de-normalize:
+    # Your training: (y - 48) / 48  => inference: y = y_pred * 48 + 48
+    # We do it safely here:
+    pred = (pred * 48.0) + 48.0
 
-    coords = keypoints_to_xy(y_pred, target_cols)
+    # Clip to valid [0, 96]
+    pred = np.clip(pred, 0.0, 96.0)
 
-    st.subheader('Predicted keypoints (first 10 values)')
-    preview_items = list(coords.items())[:10]
-    st.write({k: round(v, 3) for k, v in preview_items})
+    pts = to_xy(pred)
 
-    overlay = draw_points_on_image(gray_96, coords, point_size=2)
-    st.subheader('Keypoints overlay (96×96)')
+    st.subheader('Prediction (keypoints on 96×96)')
+    overlay = draw_keypoints(pil_img, pts)
     st.image(overlay, use_container_width=False)
 
-    # Download as JSON
-    out = {'img_size': preprocess_cfg.get('img_size', [96, 96]), 'predictions': coords}
-    st.download_button(
-        'Download predictions as JSON',
-        data=json.dumps(out, indent=2),
-        file_name='predicted_keypoints.json',
-        mime='application/json'
-    )
+    st.subheader('Keypoints table (x, y)')
+    # Build a nice table using target_cols order
+    # target_cols is typically like: ['left_eye_center_x', 'left_eye_center_y', ...]
+    rows = []
+    for i in range(0, len(target_cols), 2):
+        name_x = target_cols[i]
+        name_y = target_cols[i + 1]
+        rows.append({
+            'keypoint': name_x.replace('_x', ''),
+            'x_name': name_x,
+            'y_name': name_y,
+            'x': float(pred[i]),
+            'y': float(pred[i + 1]),
+        })
+
+    st.dataframe(rows, use_container_width=True)
+
 else:
-    st.info('Upload an image to run inference.')
+    st.info('Upload an image to get predictions.')