Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,146 +1,135 @@
 import json
-from pathlib import Path
-
 import numpy as np
 import streamlit as st
 from PIL import Image
 import tensorflow as tf
-import matplotlib.pyplot as plt
-
-
-# -------------------------
-# 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 CNN predicts 30 facial keypoint coordinates (x/y).')
+from pathlib import Path
 
+st.set_page_config(page_title='Facial Keypoints Predictor (CNN)', page_icon='🙂', layout='centered')
 
-# -------------------------
-# Paths (HF-friendly: repo root)
-# -------------------------
 BASE_DIR = Path(__file__).resolve().parent
 
-MODEL_PATH = BASE_DIR / 'final_keypoints_cnn.keras'
+# --- 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)
+
 TARGET_COLS_PATH = BASE_DIR / 'target_cols.json'
-PREPROCESS_PATH = BASE_DIR / 'preprocess_config.json'
+PREPROCESS_CFG_PATH = BASE_DIR / 'preprocess_config.json'
 
 
-# -------------------------
-# Loaders
-# -------------------------
 @st.cache_resource
-def load_model():
-    if not MODEL_PATH.exists():
+def load_assets():
+    # load metadata
+    if not TARGET_COLS_PATH.exists():
+        raise FileNotFoundError(f'Missing {TARGET_COLS_PATH.name} in repo root.')
+
+    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)
+
+    # 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(
-            f'Model not found: {MODEL_PATH.name}. Put it in the repo root (same folder as app.py).'
+            'Model not found. Upload either:\n'
+            f'- {MODEL_SAVEDMODEL_DIR.name}/ (SavedModel folder)\n'
+            f'- {MODEL_H5_PATH.name} (H5 file)\n'
         )
-    return tf.keras.models.load_model(MODEL_PATH, compile=False)
-
-
-@st.cache_data
-def load_json(path: Path):
-    if not path.exists():
-        raise FileNotFoundError(
-            f'File not found: {path.name}. Put it in the repo root (same folder as app.py).'
-        )
-    with open(path, 'r') as f:
-        return json.load(f)
-
-
-model = load_model()
-target_cols = load_json(TARGET_COLS_PATH)
-pre_cfg = load_json(PREPROCESS_PATH)
 
-IMG_H, IMG_W = pre_cfg.get('img_size', [96, 96])
+    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
+
+
+st.title('🙂 Facial Keypoints Predictor (CNN)')
+st.write('Upload a face image and the model will predict 15 facial keypoints (30 values: x/y).')
+
+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'
+    )
+
+model, target_cols, preprocess_cfg = load_assets()
 
-
-# -------------------------
-# Preprocess + Postprocess
-# -------------------------
-def preprocess_image(pil_img: Image.Image) -> np.ndarray:
-    # Convert to grayscale like training data
-    img = pil_img.convert('L')
-    img = img.resize((IMG_W, IMG_H))
-    arr = np.array(img, dtype=np.float32)  # 0..255
-    arr = arr / 255.0
-    arr = arr.reshape(1, IMG_H, IMG_W, 1)
-    return arr
-
-
-def denormalize_keypoints(pred_norm: np.ndarray) -> np.ndarray:
-    # Training normalization: (y - 48) / 48  -> invert: y = pred*48 + 48
-    pred = pred_norm * 48.0 + 48.0
-    return pred
-
-
-def plot_keypoints(pil_img: Image.Image, keypoints_xy: np.ndarray):
-    img = pil_img.convert('L').resize((IMG_W, IMG_H))
-    xs = keypoints_xy[0::2]
-    ys = keypoints_xy[1::2]
-
-    fig, ax = plt.subplots(figsize=(5, 5))
-    ax.imshow(img, cmap='gray')
-    ax.scatter(xs, ys, s=25)
-    ax.axis('off')
-    return fig
-
-
-# -------------------------
-# UI
-# -------------------------
 uploaded = st.file_uploader('Upload an image (jpg/png)', type=['jpg', 'jpeg', 'png'])
 
-with st.expander('Settings', expanded=False):
-    show_table = st.checkbox('Show predicted coordinates table', value=True)
-    show_overlay = st.checkbox('Show keypoints overlay', value=True)
-
-if uploaded is None:
-    st.info('Upload an image to get predictions.')
-    st.stop()
-
-pil_img = Image.open(uploaded)
-
-st.subheader('Input Image')
-st.image(pil_img, use_container_width=True)
-
-x = preprocess_image(pil_img)
-
-with st.spinner('Predicting keypoints...'):
-    pred_norm = model.predict(x, verbose=0)[0]  # shape (30,)
-pred_px = denormalize_keypoints(pred_norm)
+if uploaded is not None:
+    pil_img = Image.open(uploaded)
+    st.image(pil_img, caption='Uploaded image', use_container_width=True)
 
-# Build table
-rows = []
-for i, name in enumerate(target_cols):
-    rows.append({'feature': name, 'value_px': float(pred_px[i])})
+    x, gray_96 = preprocess_image(pil_img, img_size=tuple(preprocess_cfg.get('img_size', [96, 96])))
 
-if show_table:
-    st.subheader('Predicted Keypoints (Pixels)')
-    st.dataframe(rows, use_container_width=True)
+    y_pred_norm = model.predict(x, verbose=0)[0]          # (30,)
+    y_pred = denormalize_and_clip(y_pred_norm)           # (30,)
 
-if show_overlay:
-    st.subheader('Overlay')
-    fig = plot_keypoints(pil_img, pred_px)
-    st.pyplot(fig, clear_figure=True)
+    coords = keypoints_to_xy(y_pred, target_cols)
 
-# Optional: downloadable single-image "submission-like" csv (ImageId=1)
-st.subheader('Download (optional)')
-st.caption('This creates a single-image submission-like file: ImageId=1, RowId=1..30.')
+    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})
 
-csv_lines = ['RowId,ImageId,FeatureName,Location']
-for idx, (name, val) in enumerate(zip(target_cols, pred_px), start=1):
-    csv_lines.append(f'{idx},1,{name},{val}')
-csv_text = '\n'.join(csv_lines)
+    overlay = draw_points_on_image(gray_96, coords, point_size=2)
+    st.subheader('Keypoints overlay (96×96)')
+    st.image(overlay, use_container_width=False)
 
-st.download_button(
-    label='Download predicted_keypoints.csv',
-    data=csv_text.encode('utf-8'),
-    file_name='predicted_keypoints.csv',
-    mime='text/csv'
-)
+    # 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'
+    )
+else:
+    st.info('Upload an image to run inference.')