Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,30 @@
+.vscode/
+.vs/
+.venv/
+venv/
+__pycache__/
+*.py[cod]
+*.egg-info/
+dist/
+build/
+*.h5
+*.keras
+saved_model_*/
+*.pb
+checkpoints/
+*.ckpt
+data/
+data/**
+UTKFace/
+models/
+models/**
+*.log
+*.npy
+*.npz
+*.tflite
+*.zip
+*.tar.gz
+.env
+.env.*
+.DS_Store
+.ipynb_checkpoints/

README.md

@@ -1,3 +1,168 @@
----
-license: mit
----
+---
+language: en
+license: mit
+tags: ["image-regression", "tensorflow", "mobilenetv2", "utkface", "age-estimation"]
+datasets: ["UTKFace"]
+metrics: ["mean_absolute_error"]
+---
+
+# UTKFace Age Regression — Model Card
+
+This repository contains code to train a TensorFlow / Keras regression model that estimates a person's age from a face image using the UTKFace dataset. The model uses a MobileNetV2 backbone and a small regression head on top.
+
+## Summary
+
+- **Model type**: Image regression (single-output continuous)
+- **Backbone**: MobileNetV2 (ImageNet pre-trained)
+- **Task**: Age estimation (years)
+- **Dataset**: UTKFace (public dataset; filenames encode age)
+- **Reported metric**: Mean Absolute Error (MAE) — see Evaluation section for how to compute and report MAE for your runs
+
+## Model details
+
+- **Input**: RGB face image (recommended size: 224×224)
+- **Output**: Single scalar value — predicted age in years
+- **Preprocessing**: MobileNetV2 preprocessing (scales inputs to [-1, 1])
+- **Loss**: Mean Squared Error (MSE) used during training
+- **Metric for reporting**: Mean Absolute Error (MAE)
+
+## Intended uses
+
+- Research and educational purposes for learning about image regression and age estimation
+- Prototyping demo applications that predict approximate age ranges from face crops
+
+## Out-of-scope / Limitations
+
+- This model provides an estimate of age; it's not a substitute for official identification
+- Models trained on UTKFace carry dataset biases (race, gender, age distribution). They may underperform on underrepresented groups.
+- Do not use this model for high-stakes decision making (employment, legal, medical, etc.)
+
+## Dataset
+
+**UTKFace**
+
+- **Source**: https://susanqq.github.io/UTKFace/
+- **Format**: Filenames encode metadata as `<age>_<gender>_<race>_<date&time>.jpg`.
+- **Usage**: The training scripts in this repo extract the age from the filename (the integer before the first underscore).
+- **Note**: Respect the dataset's license and authors when redistributing or publishing results.
+
+## Training details
+
+- **Framework**: TensorFlow / Keras
+- **Backbone**: MobileNetV2 pretrained on ImageNet
+- **Head**: GlobalAveragePooling2D -> Dense(128, relu) -> Dense(1, linear)
+- **Recommended input size**: 224×224 (configurable via command-line args in `train.py`)
+- **Batch size**: configurable (default set in `train.py`)
+- **Optimizer**: Adam (default), learning rate and scheduler configurable in `train.py`
+- **Loss**: Mean Squared Error (MSE)
+- **Metric**: Mean Absolute Error (MAE) reported on validation/test sets
+- **Augmentations**: Basic augmentations recommended (flip, random crop/brightness) for better robustness
+
+## Reproducibility / Example training command
+
+1. **Prepare UTKFace dataset**
+   - Download and extract UTKFace images into `data/UTKFace/` or pass `--dataset_dir` to the training script.
+2. **Install dependencies**
+   - `python -m pip install -r requirements.txt`
+3. **Train**
+   - `python train.py --dataset_dir data/UTKFace --epochs 30 --batch_size 32 --img_size 224 --output_dir saved_model`
+
+The `train.py` script builds a tf.data pipeline, extracts ages from filenames, constructs a MobileNetV2-based model, and saves the trained model to the `--output_dir`.
+
+## Evaluation and metrics (MAE)
+
+Mean Absolute Error (MAE) gives an intuitive measure of average error in predicted age (in years):
+
+```
+MAE = mean(|y_true - y_pred|)
+```
+
+Compute MAE in Python (example):
+
+```python
+import numpy as np
+mae = np.mean(np.abs(y_true - y_pred))
+```
+
+Example: the training script prints per-epoch validation MAE. To reproduce test MAE after training, run the provided evaluation routine or:
+
+```python
+from tensorflow import keras
+import numpy as np
+model = keras.models.load_model('saved_model')
+# prepare test_images, test_labels arrays
+preds = model.predict(test_images).squeeze()
+mae = float(np.mean(np.abs(test_labels - preds)))
+print('Test MAE (years):', mae)
+```
+
+Note: Exact MAE depends on preprocessing, train/validation split, augmentations, and hyperparameters. Report MAE alongside the exact training configuration for reproducibility.
+
+## Usage — Quick examples
+
+**Python (local SavedModel)**
+
+```python
+import tensorflow as tf
+import numpy as np
+from PIL import Image
+from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
+
+model = tf.keras.models.load_model('saved_model')  # path to a SavedModel directory
+img = Image.open('path/to/face.jpg').convert('RGB').resize((224, 224))
+arr = np.array(img, dtype=np.float32)
+arr = preprocess_input(arr)
+pred = model.predict(np.expand_dims(arr, 0))[0, 0]
+print('Predicted age (years):', float(pred))
+```
+
+**Command-line (using predict.py)**
+
+```
+python predict.py --model_dir saved_model --image path/to/face.jpg
+```
+
+**Loading from Hugging Face Hub**
+
+If you upload your saved model to the Hugging Face Hub, Consumers can download it using the `huggingface_hub` package. For example, in a Space, set the environment variable `HF_MODEL_ID` to the model repository (e.g. `username/my-age-model`) and the Gradio app supplied in this repo will attempt to download and use it.
+
+**Gradio demo / Hugging Face Space**
+
+A simple Gradio app is provided in `app.py` that:
+
+- accepts an input face image
+- preprocesses it (224×224 + MobileNetV2 preprocess)
+- returns the predicted age (years) and the model's raw output
+
+**How to host as a Space**
+
+1. Create a new Space on Hugging Face and select "Gradio" as the SDK.
+2. Push this repository to the Space (include `app.py`, your `saved_model/` directory or set `HF_MODEL_ID` to your model on the Hub).
+3. Make sure `requirements.txt` includes `gradio` and `huggingface_hub` (the repository `requirements.txt` in this project may be extended with these packages for the Space).
+
+## Files in this repository
+
+- `train.py` — training script
+- `predict.py` — single-image prediction helper
+- `convert_model.py` — conversion helpers
+- `inference_log.py`, `inference_log.txt`, `load_predict_log.txt` — logging and CLI helpers for inference (dev)
+- `app.py` — (added) Gradio demo app for live predictions
+- `requirements.txt` — Python dependencies (extend for Spaces with `gradio` and `huggingface_hub`)
+
+## Security, biases and ethical considerations
+
+- Age estimation models can reflect and amplify biases in the training data (race and gender imbalance, age distribution). Evaluate fairness across demographic slices before using widely.
+- Avoid using the model in high-risk contexts where inaccurate age estimates could cause harm.
+
+## How to cite / license
+
+- UTKFace authors and dataset should be cited if you publish results.
+- This repository is provided under the MIT license (see LICENSE file if present).
+
+## Contact and credits
+
+**Maintainer**: Stealth Labs Ltd.
+
+**Acknowledgements**
+
+Thanks to the UTKFace dataset authors for the publicly available images used in training and experimentation.

app.py

@@ -0,0 +1,74 @@
+import os
+import numpy as np
+from PIL import Image
+import tensorflow as tf
+from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
+import gradio as gr
+
+# Try to load a local SavedModel first, otherwise try to download from Hugging Face Hub
+MODEL_DIR = "saved_model"
+model = None
+
+if os.path.isdir(MODEL_DIR):
+    try:
+        model = tf.keras.models.load_model(MODEL_DIR)
+        print(f"Loaded model from local path: {MODEL_DIR}")
+    except Exception as e:
+        print(f"Failed to load local model: {e}")
+
+if model is None:
+    # If HF_MODEL_ID is set, attempt to download model files from the Hub
+    HF_MODEL_ID = os.environ.get("HF_MODEL_ID")
+    if HF_MODEL_ID:
+        try:
+            from huggingface_hub import snapshot_download
+            repo_dir = snapshot_download(repo_id=HF_MODEL_ID)
+            # Expecting a SavedModel dir inside the repo; try to load from repo root
+            model = tf.keras.models.load_model(repo_dir)
+            print(f"Loaded model from HF Hub repo: {HF_MODEL_ID}")
+        except Exception as e:
+            print(f"Failed to load model from HF Hub ({HF_MODEL_ID}): {e}")
+
+if model is None:
+    raise RuntimeError(
+        "No model found. Place a SavedModel in './saved_model' or set HF_MODEL_ID env var to a Hugging Face model repo containing a SavedModel."
+    )
+
+INPUT_SIZE = (224, 224)
+
+
+def predict_age(image: Image.Image):
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+    image = image.resize(INPUT_SIZE)
+    arr = np.array(image).astype(np.float32)
+    arr = preprocess_input(arr)
+    arr = np.expand_dims(arr, 0)
+
+    pred = model.predict(arr)[0]
+    # Ensure scalar
+    if hasattr(pred, '__len__'):
+        pred = float(np.asarray(pred).squeeze())
+    else:
+        pred = float(pred)
+
+    return {
+        "predicted_age": round(pred, 2),
+        "raw_output": float(pred)
+    }
+
+
+demo = gr.Interface(
+    fn=predict_age,
+    inputs=gr.Image(type='pil', label='Face image (crop to face for best results)'),
+    outputs=[
+        gr.Number(label='Predicted age (years)'),
+        gr.Number(label='Raw model output')
+    ],
+    examples=[],
+    title='UTKFace Age Estimator',
+    description='Upload a cropped face image and the model will predict age in years. For Spaces, set the HF_MODEL_ID environment variable to your Hugging Face model repo if you want the app to download a SavedModel from the Hub.'
+)
+
+if __name__ == '__main__':
+    demo.launch(server_name='0.0.0.0', server_port=int(os.environ.get('PORT', 7860)))

convert_model.py

@@ -0,0 +1,27 @@
+import tensorflow as tf
+
+print('loading best_model.h5...')
+try:
+    # Load without compiling to avoid deserializing legacy training configs/metrics
+    m = tf.keras.models.load_model('best_model.h5', compile=False)
+except Exception as e:
+    print('Failed to load best_model.h5:', e)
+    raise
+
+# Try to export to the TF SavedModel format first
+try:
+    m.export('saved_model_age_regressor')
+    print('Exported SavedModel to ./saved_model_age_regressor')
+except Exception as e:
+    print('Export to SavedModel failed:', e)
+    # Fallback: save as Keras native single-file and HDF5 for compatibility
+    try:
+        m.save('saved_model_age_regressor.keras')
+        print('Saved Keras model to ./saved_model_age_regressor.keras')
+    except Exception as e2:
+        print('Saving Keras native format failed:', e2)
+    try:
+        m.save('final_model.h5')
+        print('Saved HDF5 model to ./final_model.h5')
+    except Exception as e3:
+        print('Saving HDF5 format failed:', e3)

inference_log.py

@@ -0,0 +1,47 @@
+import traceback
+from pathlib import Path
+
+log_path = Path('inference_log.txt')
+with log_path.open('w', encoding='utf-8') as f:
+    def log(*args, **kwargs):
+        print(*args, file=f, **kwargs)
+        f.flush()
+
+    try:
+        log('Starting inference log')
+        import tensorflow as tf
+        import numpy as np
+        from PIL import Image
+
+        model_path = 'saved_model_age_regressor'
+        img_path = Path('data/UTKFace/53_1_1_20170110122449716.jpg.chip.jpg')
+
+        log('Model path:', model_path)
+        log('Image path:', str(img_path))
+
+        log('Attempting to load model with compile=False...')
+        m = tf.keras.models.load_model(model_path, compile=False)
+        log('Loaded model type:', type(m))
+
+        try:
+            m.summary(print_fn=lambda *a, **k: log(*a, **k))
+        except Exception as e:
+            log('model.summary failed:', e)
+
+        img = Image.open(img_path).convert('RGB').resize((224,224))
+        arr = np.array(img, dtype=np.float32)/255.0
+        x = np.expand_dims(arr, 0)
+        log('Input shape:', x.shape)
+
+        log('Running predict...')
+        pred = m.predict(x)
+        log('Raw prediction output:', pred, 'shape:', getattr(pred, 'shape', None))
+        try:
+            log('Predicted age:', float(pred.flatten()[0]))
+        except Exception as e:
+            log('Error converting prediction to float:', e)
+
+        log('Inference finished successfully')
+    except Exception:
+        traceback.print_exc(file=f)
+        log('Inference script caught exception')

inference_log.txt

@@ -0,0 +1,12 @@
+Starting inference log
+Model path: saved_model_age_regressor
+Image path: data\UTKFace\53_1_1_20170110122449716.jpg.chip.jpg
+Attempting to load model with compile=False...
+Traceback (most recent call last):
+  File "C:\Users\SammyHarris\Downloads\Age-classification-model\inference_log.py", line 23, in <module>
+    m = tf.keras.models.load_model(model_path, compile=False)
+        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "C:\Users\SammyHarris\Downloads\Age-classification-model\.venv\Lib\site-packages\keras\src\saving\saving_api.py", line 209, in load_model
+    raise ValueError(
+ValueError: File format not supported: filepath=saved_model_age_regressor. Keras 3 only supports V3 `.keras` files and legacy H5 format files (`.h5` extension). Note that the legacy SavedModel format is not supported by `load_model()` in Keras 3. In order to reload a TensorFlow SavedModel as an inference-only layer in Keras 3, use `keras.layers.TFSMLayer(saved_model_age_regressor, call_endpoint='serving_default')` (note that your `call_endpoint` might have a different name).
+Inference script caught exception

load_predict_log.txt

@@ -0,0 +1,7 @@
+Starting model load & predict test
+Image path: data\UTKFace\53_1_1_20170110122449716.jpg.chip.jpg
+HDF5/.keras not found; attempting to wrap TF SavedModel using TFSMLayer...
+Wrapper model created; running predict...
+Prediction returned a dict with keys: ['output_0']
+Output 'output_0': shape=(1, 1) values=[28.5549259185791]
+Finished load & predict test

predict.py

@@ -0,0 +1,88 @@
+"""
+Load a trained age regression model and run a prediction on a single image.
+
+Usage: python predict.py --model_path saved_model_age_regressor --image_path some_image.jpg
+"""
+import argparse
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+import tensorflow as tf
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model_path', type=str, default='saved_model_age_regressor')
+    parser.add_argument('--image_path', type=str, required=True)
+    parser.add_argument('--img_size', type=int, default=224)
+    parser.add_argument('--output_key', type=str, default=None,
+                        help='If the model returns a dict, select this key for the numeric prediction. If omitted the first numeric output will be used.')
+    return parser.parse_args()
+
+
+def load_image(path, img_size):
+    img = Image.open(path).convert('RGB')
+    img = img.resize((img_size, img_size))
+    arr = np.array(img, dtype=np.float32) / 255.0
+    return arr
+
+
+def main():
+    args = parse_args()
+    model_path = Path(args.model_path)
+    # Load Keras .h5/.keras files directly, and attempt Keras load for directories first.
+    if model_path.is_file() and model_path.suffix.lower() in ('.h5', '.keras'):
+        model = tf.keras.models.load_model(str(model_path), compile=False)
+        print(f"Loaded Keras model file: {model_path}")
+    elif model_path.is_dir():
+        # Some SavedModel directories are not loadable with tf.keras.load_model in Keras 3;
+        # try load_model first (covers .keras saved dirs), otherwise wrap with TFSMLayer.
+        try:
+            model = tf.keras.models.load_model(str(model_path), compile=False)
+            print(f"Loaded Keras-compatible model from directory: {model_path}")
+        except Exception:
+            # Wrap the SavedModel with a TFSMLayer for inference compatibility in Keras.
+            try:
+                tf_layer = tf.keras.layers.TFSMLayer(str(model_path), call_endpoint='serving_default')
+                model = tf.keras.Sequential([
+                    tf.keras.Input(shape=(args.img_size, args.img_size, 3)),
+                    tf_layer,
+                ])
+                print(f"Wrapped TensorFlow SavedModel at {model_path} with TFSMLayer (serving_default).")
+            except Exception as e:
+                raise RuntimeError(f"Failed to load or wrap SavedModel directory '{model_path}': {e}")
+    else:
+        # Unknown path type: try load_model and allow it to raise a helpful exception.
+        model = tf.keras.models.load_model(str(model_path), compile=False)
+        print(f"Loaded model from path: {model_path}")
+    image_path = Path(args.image_path)
+    if not image_path.exists():
+        raise FileNotFoundError(f"Image not found: {image_path}")
+    x = load_image(image_path, args.img_size)
+    x = np.expand_dims(x, axis=0)
+    pred = model.predict(x)
+
+    # If the model returns a dict (typical for a wrapped SavedModel serving signature),
+    # select the requested output key or fall back to the first available numeric output.
+    if isinstance(pred, dict):
+        if args.output_key:
+            if args.output_key not in pred:
+                raise KeyError(f"Requested output key '{args.output_key}' not found. Available keys: {list(pred.keys())}")
+            chosen = pred[args.output_key]
+        else:
+            first_key = next(iter(pred.keys()))
+            print(f"No --output_key provided; using first output key: '{first_key}'")
+            chosen = pred[first_key]
+        arr = np.asarray(chosen)
+    else:
+        arr = np.asarray(pred)
+
+    if arr.size == 0:
+        raise ValueError("Model returned an empty prediction.")
+    age_pred = float(arr.flatten()[0])
+    print(f"Predicted age: {age_pred:.2f} years")
+
+
+if __name__ == '__main__':
+    main()

requirements.txt

@@ -0,0 +1,6 @@
+tensorflow>=2.10.0
+numpy
+pillow
+matplotlib
+requests
+tqdm

run_load_and_predict.py

@@ -0,0 +1,91 @@
+import traceback
+from pathlib import Path
+
+log_path = Path('load_predict_log.txt')
+with log_path.open('w', encoding='utf-8') as f:
+    def log(*args, **kwargs):
+        print(*args, file=f, **kwargs)
+        f.flush()
+
+    try:
+        log('Starting model load & predict test')
+        import tensorflow as tf
+        import numpy as np
+        from PIL import Image
+        import os
+
+        img_path = Path('data/UTKFace/53_1_1_20170110122449716.jpg.chip.jpg')
+        log('Image path:', str(img_path))
+
+        # Try HDF5 / .h5 first
+        h5_path = Path('final_model.h5')
+        keras_path = Path('saved_model_age_regressor.keras')
+        saved_model_dir = Path('saved_model_age_regressor')
+
+        if h5_path.exists():
+            try:
+                log('Attempting to load HDF5 model:', str(h5_path))
+                m = tf.keras.models.load_model(str(h5_path), compile=False)
+                log('Loaded HDF5 model:', type(m))
+                img = Image.open(img_path).convert('RGB').resize((224,224))
+                x = np.expand_dims(np.array(img, dtype=np.float32)/255.0, 0)
+                log('Running predict on HDF5 model...')
+                pred = m.predict(x)
+                log('Prediction result (HDF5):', pred.tolist())
+            except Exception:
+                log('Exception while loading/predicting from HDF5:')
+                traceback.print_exc(file=f)
+        elif keras_path.exists():
+            try:
+                log('Attempting to load Keras native file:', str(keras_path))
+                m = tf.keras.models.load_model(str(keras_path), compile=False)
+                log('Loaded Keras native model:', type(m))
+                img = Image.open(img_path).convert('RGB').resize((224,224))
+                x = np.expand_dims(np.array(img, dtype=np.float32)/255.0, 0)
+                pred = m.predict(x)
+                log('Prediction result (KERAS):', pred.tolist())
+            except Exception:
+                log('Exception while loading/predicting from Keras file:')
+                traceback.print_exc(file=f)
+        elif saved_model_dir.exists():
+            try:
+                log('HDF5/.keras not found; attempting to wrap TF SavedModel using TFSMLayer...')
+                try:
+                    from keras.layers import TFSMLayer
+                except Exception as e:
+                    log('TFSMLayer import failed:', e)
+                    raise
+                # Build wrapper model
+                inputs = tf.keras.Input(shape=(224,224,3))
+                tfsml = TFSMLayer(str(saved_model_dir), call_endpoint='serving_default')
+                outputs = tfsml(inputs)
+                wrapper = tf.keras.Model(inputs, outputs)
+                log('Wrapper model created; running predict...')
+                img = Image.open(img_path).convert('RGB').resize((224,224))
+                x = np.expand_dims(np.array(img, dtype=np.float32)/255.0, 0)
+                pred = wrapper.predict(x)
+                # The SavedModel serving signature can return a dict mapping names->arrays
+                if isinstance(pred, dict):
+                    log('Prediction returned a dict with keys:', list(pred.keys()))
+                    import numpy as _np
+                    for k, v in pred.items():
+                        try:
+                            arr = _np.array(v)
+                            log(f"Output '{k}': shape={arr.shape} values={arr.flatten()[:10].tolist()}")
+                        except Exception as _e:
+                            log(f"Could not convert output '{k}' to numpy array:", _e)
+                else:
+                    try:
+                        log('Prediction result (wrapped SavedModel):', pred.tolist())
+                    except Exception:
+                        log('Prediction result (wrapped SavedModel) type:', type(pred))
+            except Exception:
+                log('Exception while wrapping/using SavedModel:')
+                traceback.print_exc(file=f)
+        else:
+            log('No model file found: looked for final_model.h5, saved_model_age_regressor.keras, or saved_model_age_regressor/')
+
+        log('Finished load & predict test')
+    except Exception:
+        traceback.print_exc(file=f)
+        log('Top-level exception')

space_requirements.txt

@@ -0,0 +1,5 @@
+gradio
+huggingface_hub
+tensorflow>=2.10.0
+pillow
+numpy

test_predict.py

@@ -0,0 +1,27 @@
+import tensorflow as tf
+from pathlib import Path
+import numpy as np
+from PIL import Image
+
+model_path = 'saved_model_age_regressor'
+img_path = Path('data/UTKFace/53_1_1_20170110122449716.jpg.chip.jpg')
+print('Model path:', model_path, flush=True)
+print('Image path:', img_path, flush=True)
+
+m = tf.keras.models.load_model(model_path, compile=False)
+print('Loaded model type:', type(m), flush=True)
+try:
+    m.summary()
+except Exception as e:
+    print('model.summary failed:', e, flush=True)
+
+img = Image.open(img_path).convert('RGB').resize((224,224))
+arr = np.array(img, dtype=np.float32)/255.0
+x = np.expand_dims(arr, 0)
+print('Input shape:', x.shape, flush=True)
+pred = m.predict(x)
+print('Raw prediction output:', pred, 'shape:', getattr(pred, 'shape', None), flush=True)
+try:
+    print('Predicted age:', float(pred.flatten()[0]), flush=True)
+except Exception as e:
+    print('Error converting prediction to float:', e, flush=True)

train.py

@@ -0,0 +1,273 @@
+"""
+Train a TensorFlow regression model to predict age from face images (UTKFace dataset).
+
+Usage:
+  - Put UTKFace images into a folder, e.g. data/UTKFace/
+  - python train.py --dataset_dir data/UTKFace --epochs 30 --batch_size 32
+
+The script extracts the age from the filename (before the first underscore).
+"""
+
+import os
+import argparse
+import random
+import math
+import zipfile
+from pathlib import Path
+
+import numpy as np
+from tqdm import tqdm
+import requests
+
+import tensorflow as tf
+from tensorflow import keras
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Train an age regression model on UTKFace images")
+    parser.add_argument("--dataset_dir", type=str, default="data/UTKFace", help="Path to folder containing UTKFace images")
+    parser.add_argument("--img_size", type=int, default=224, help="Image size (square)")
+    parser.add_argument("--batch_size", type=int, default=32)
+    parser.add_argument("--epochs", type=int, default=30)
+    parser.add_argument("--val_split", type=float, default=0.12, help="Fraction to reserve for validation")
+    parser.add_argument("--learning_rate", type=float, default=1e-4)
+    parser.add_argument("--auto_download", type=lambda x: (str(x).lower() in ("true", "1", "yes")), default=False,
+                        help="Whether to attempt to download UTKFace archive automatically if dataset folder is missing")
+    parser.add_argument("--fine_tune", type=lambda x: (str(x).lower() in ("true", "1", "yes")), default=False,
+                        help="Whether to unfreeze part of the backbone for fine-tuning")
+    args = parser.parse_args()
+    return args
+
+
+def attempt_download_utkface(dest_dir: Path):
+    """Attempt to download a ZIP archive of the UTKFace repository and extract it.
+
+    This may fail if the remote hosting changes. The function attempts a best-effort download
+    from the repository URL commonly used to host UTKFace on GitHub.
+    """
+    dest_dir.mkdir(parents=True, exist_ok=True)
+    github_zip = "https://github.com/susanqq/UTKFace/archive/refs/heads/master.zip"
+    tmp_zip = dest_dir / "utkface_master.zip"
+    print(f"Attempting to download UTKFace from {github_zip} ...")
+
+    try:
+        with requests.get(github_zip, stream=True, timeout=30) as r:
+            r.raise_for_status()
+            total = int(r.headers.get('content-length', 0))
+            with open(tmp_zip, 'wb') as f:
+                for chunk in r.iter_content(chunk_size=8192):
+                    if chunk:
+                        f.write(chunk)
+
+        print("Download complete. Extracting archive...")
+        with zipfile.ZipFile(tmp_zip, 'r') as z:
+            z.extractall(dest_dir)
+
+        # Move images into dest_dir if they're inside a top-level folder
+        extracted_root = None
+        for name in os.listdir(dest_dir):
+            if name.lower().startswith('utkface') and os.path.isdir(dest_dir / name):
+                extracted_root = dest_dir / name
+                break
+        if extracted_root:
+            images = list(extracted_root.rglob('*.jpg')) + list(extracted_root.rglob('*.png'))
+            for p in images:
+                target = dest_dir / p.name
+                try:
+                    os.replace(p, target)
+                except Exception:
+                    pass
+        # clean up
+        try:
+            os.remove(tmp_zip)
+        except Exception:
+            pass
+        print("UTKFace images should now be in:", dest_dir)
+    except Exception as e:
+        print("Automatic download failed:", e)
+        print("Please download the UTKFace archive manually and place images in the dataset directory.")
+
+
+def collect_image_paths_and_labels(dataset_dir: Path):
+    # UTKFace filenames: <age>_<gender>_<race>_<date&time>.jpg
+    img_paths = []
+    labels = []
+    supported_ext = ('.jpg', '.jpeg', '.png')
+    for p in dataset_dir.iterdir():
+        if p.is_file() and p.suffix.lower() in supported_ext:
+            # parse age
+            parts = p.name.split('_')
+            try:
+                age = int(parts[0])
+            except Exception:
+                continue
+            img_paths.append(str(p))
+            labels.append(age)
+    return img_paths, labels
+
+
+def make_dataset(paths, labels, img_size, batch_size, is_training=True):
+    paths = tf.convert_to_tensor(paths)
+    labels = tf.convert_to_tensor(labels, dtype=tf.float32)
+
+    ds = tf.data.Dataset.from_tensor_slices((paths, labels))
+    if is_training:
+        ds = ds.shuffle(10000, reshuffle_each_iteration=True)
+
+    def _load_image(path, label):
+        img = tf.io.read_file(path)
+        img = tf.image.decode_jpeg(img, channels=3)
+        img = tf.image.resize(img, [img_size, img_size])
+        img = img / 255.0  # normalize to [0,1]
+        if is_training:
+            img = data_augmentation(img)
+        return img, label
+
+    ds = ds.map(_load_image, num_parallel_calls=tf.data.AUTOTUNE)
+    ds = ds.batch(batch_size).prefetch(tf.data.AUTOTUNE)
+    return ds
+
+
+def data_augmentation(image):
+    # Simple augmentation pipeline
+    image = tf.image.random_flip_left_right(image)
+    image = tf.image.random_brightness(image, max_delta=0.08)
+    image = tf.image.random_contrast(image, 0.9, 1.1)
+    # random zoom by central crop/resizing
+    if tf.random.uniform(()) > 0.6:
+        crop_frac = tf.random.uniform((), 0.8, 1.0)
+        shape = tf.shape(image)
+        crop_h = tf.cast(tf.cast(shape[0], tf.float32) * crop_frac, tf.int32)
+        crop_w = tf.cast(tf.cast(shape[1], tf.float32) * crop_frac, tf.int32)
+        image = tf.image.random_crop(image, size=[crop_h, crop_w, 3])
+        image = tf.image.resize(image, [shape[0], shape[1]])
+    return image
+
+
+def build_model(img_size, fine_tune=False):
+    inputs = keras.Input(shape=(img_size, img_size, 3))
+    base = keras.applications.MobileNetV2(include_top=False, input_tensor=inputs, weights='imagenet')
+    base.trainable = False
+
+    x = base.output
+    x = keras.layers.GlobalAveragePooling2D()(x)
+    x = keras.layers.Dropout(0.2)(x)
+    x = keras.layers.Dense(128, activation='relu')(x)
+    x = keras.layers.Dense(64, activation='relu')(x)
+    outputs = keras.layers.Dense(1, name='age')(x)  # regression output
+
+    model = keras.Model(inputs=inputs, outputs=outputs)
+
+    if fine_tune:
+        # Unfreeze last blocks for fine-tuning
+        base.trainable = True
+        # Freeze earlier layers
+        for layer in base.layers[:-30]:
+            layer.trainable = False
+
+    return model
+
+
+def main():
+    args = parse_args()
+    dataset_dir = Path(args.dataset_dir)
+
+    if (not dataset_dir.exists() or not any(dataset_dir.iterdir())) and args.auto_download:
+        attempt_download_utkface(dataset_dir)
+
+    if not dataset_dir.exists() or not any(dataset_dir.iterdir()):
+        raise RuntimeError(f"No images found in {dataset_dir}. Place UTKFace images there or use --auto_download True to attempt download.")
+
+    paths, labels = collect_image_paths_and_labels(dataset_dir)
+    if len(paths) == 0:
+        raise RuntimeError("No valid UTKFace images found in dataset directory. Ensure the files follow the naming convention '<age>_...'.")
+
+    # Convert to numpy lists
+    paths = np.array(paths)
+    labels = np.array(labels, dtype=np.float32)
+
+    # Shuffle and split
+    indices = np.arange(len(paths))
+    np.random.shuffle(indices)
+    paths = paths[indices]
+    labels = labels[indices]
+
+    n_val = max(1, int(len(paths) * args.val_split))
+    val_paths = paths[:n_val].tolist()
+    val_labels = labels[:n_val].tolist()
+    train_paths = paths[n_val:].tolist()
+    train_labels = labels[n_val:].tolist()
+
+    print(f"Found {len(train_paths)} training images and {len(val_paths)} validation images.")
+
+    train_ds = make_dataset(train_paths, train_labels, args.img_size, args.batch_size, is_training=True)
+    val_ds = make_dataset(val_paths, val_labels, args.img_size, args.batch_size, is_training=False)
+
+    model = build_model(args.img_size, fine_tune=args.fine_tune)
+    model.compile(optimizer=keras.optimizers.Adam(learning_rate=args.learning_rate),
+                  loss='mse',
+                  metrics=[keras.metrics.MeanAbsoluteError(name='mae')])
+
+    model.summary()
+
+    callbacks = [
+        keras.callbacks.ModelCheckpoint('best_model.h5', save_best_only=True, monitor='val_loss'),
+        keras.callbacks.EarlyStopping(monitor='val_loss', patience=8, restore_best_weights=True),
+        keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=4, min_lr=1e-7)
+    ]
+
+    history = model.fit(train_ds, validation_data=val_ds, epochs=args.epochs, callbacks=callbacks)
+
+    # Evaluate
+    print("Evaluating on validation set:")
+    eval_res = model.evaluate(val_ds)
+    print(dict(zip(model.metrics_names, eval_res)))
+
+    # Save in both SavedModel (preferred) and Keras formats for compatibility
+    try:
+        # Preferred: export to SavedModel directory for TFServing/TFLite
+        model.export('saved_model_age_regressor')
+        print('Exported SavedModel to ./saved_model_age_regressor')
+    except Exception as e:
+        print('SavedModel export failed:', e)
+        # Fallback: save as Keras native single-file (.keras)
+        try:
+            model.save('saved_model_age_regressor.keras')
+            print('Saved Keras model to ./saved_model_age_regressor.keras')
+        except Exception as e2:
+            print('Keras native save failed:', e2)
+    # Also save an HDF5 copy for backward compatibility with tools that require .h5
+    try:
+        model.save('final_model.h5')
+        print('Saved HDF5 model to ./final_model.h5')
+    except Exception as e3:
+        print('HDF5 save failed:', e3)
+
+    # Show a few sample predictions
+    sample_paths = val_paths[:12]
+    sample_labels = val_labels[:12]
+
+    sample_ds = make_dataset(sample_paths, sample_labels, args.img_size, batch_size=12, is_training=False)
+    imgs, labs = next(iter(sample_ds))
+    preds = model.predict(imgs).flatten()
+
+    try:
+        import matplotlib.pyplot as plt
+        n = len(preds)
+        cols = 4
+        rows = math.ceil(n / cols)
+        plt.figure(figsize=(cols * 3, rows * 3))
+        for i in range(n):
+            ax = plt.subplot(rows, cols, i + 1)
+            img = imgs[i].numpy()
+            plt.imshow(img)
+            plt.axis('off')
+            plt.title(f"True: {int(labs[i])}\nPred: {preds[i]:.1f}")
+        plt.tight_layout()
+        plt.show()
+    except Exception:
+        print("Matplotlib not available or running headless; skipping sample visualization.")
+
+
+if __name__ == '__main__':
+    main()