Upload 129 files

.gitattributes

@@ -33,3 +33,18 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+SlimFace/src/slimface/models/detection_models/mtcnn_pytorch/caffe_models/det2.caffemodel filter=lfs diff=lfs merge=lfs -text
+SlimFace/src/slimface/models/detection_models/mtcnn_pytorch/caffe_models/det3.caffemodel filter=lfs diff=lfs merge=lfs -text
+SlimFace/src/slimface/models/detection_models/mtcnn_pytorch/caffe_models/det4.caffemodel filter=lfs diff=lfs merge=lfs -text
+SlimFace/src/slimface/models/detection_models/mtcnn_pytorch/images/example.png filter=lfs diff=lfs merge=lfs -text
+SlimFace/src/slimface/models/detection_models/mtcnn_pytorch/images/office2.jpg filter=lfs diff=lfs merge=lfs -text
+SlimFace/src/slimface/models/detection_models/mtcnn_pytorch/images/office4.jpg filter=lfs diff=lfs merge=lfs -text
+SlimFace/tests/test_images/Cate[[:space:]]Blanchett.jpg filter=lfs diff=lfs merge=lfs -text
+SlimFace/tests/test_images/Daniel[[:space:]]Day-Lewis.jpg filter=lfs diff=lfs merge=lfs -text
+SlimFace/tests/test_images/dont_know.jpg filter=lfs diff=lfs merge=lfs -text
+SlimFace/tests/test_images/Elon_Musk.jpg filter=lfs diff=lfs merge=lfs -text
+SlimFace/tests/test_images/Gal[[:space:]]Gado.jpg filter=lfs diff=lfs merge=lfs -text
+SlimFace/tests/test_images/Kate[[:space:]]Winslet.jpg filter=lfs diff=lfs merge=lfs -text
+SlimFace/tests/test_images/Tom[[:space:]]Cruise.jpg filter=lfs diff=lfs merge=lfs -text
+SlimFace/tests/test_images/Tom[[:space:]]Hanks.jpg filter=lfs diff=lfs merge=lfs -text
+SlimFace/tests/test_images/Viola[[:space:]]Davis.jpg filter=lfs diff=lfs merge=lfs -text

SlimFace/.gitignore

@@ -0,0 +1,197 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Abstra
+# Abstra is an AI-powered process automation framework.
+# Ignore directories containing user credentials, local state, and settings.
+# Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#  you could uncomment the following to ignore the enitre vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Cursor
+#  Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
+#  exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
+#  refer to https://docs.cursor.com/context/ignore-files
+.cursorignore
+.cursorindexingignore
+
+# lightning_logs
+lightning_logs

SlimFace/.python-version

@@ -0,0 +1 @@
+3.11.11

SlimFace/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Danh Tran
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

SlimFace/README.md

@@ -0,0 +1,93 @@
+# SlimFace: Slim Face Recognition
+
+> ## Credits and Citation
+>
+> ℹ️ This project is based on the [![Built on edgeface](https://img.shields.io/badge/Built%20on-otroshi%2Fedgeface-blue?style=flat&logo=github)](https://github.com/otroshi/edgeface) by [![Hatef Otroshi](https://img.shields.io/badge/GitHub-Hatef_Otroshi-blue?style=flat&logo=github)](https://github.com/otroshi), and includes our own bug fixes and enhancements.
+>
+> If this project is helpful for your research, please consider citing the original paper:
+>
+> **Edgeface: Efficient face recognition model for edge devices**  
+> *George, Anjith and Ecabert, Christophe and Shahreza, Hatef Otroshi and Kotwal, Ketan and Marcel, Sebastien*  
+> *IEEE Transactions on Biometrics, Behavior, and Identity Science (2024)*
+>
+> **If you use this work in your research, please cite the original paper:**
+> ```bibtex
+> @article{edgeface,
+>   title={Edgeface: Efficient face recognition model for edge devices},
+>   author={George, Anjith and Ecabert, Christophe and Shahreza, Hatef Otroshi and Kotwal, Ketan and Marcel, Sebastien},
+>   journal={IEEE Transactions on Biometrics, Behavior, and Identity Science},
+>   year={2024}
+> }
+> ```
+
+
+## Usage
+### Clone Repositories
+```bash
+# Clone the repository
+git clone https://github.com/danhtran2mind/SlimFace
+
+# Navigate into the newly created 'slimface' directory.
+cd SlimFace
+```
+### Install Dependencies
+**If Open-CV (CV2) does not work, run below CLI**
+```bash
+sudo apt update
+sudo apt install -y libglib2.0-0
+sudo apt install -y libgl1-mesa-dev
+```
+### Default install Dependencies
+```bash
+pip install -r requirements/requirements.txt
+```
+### Other install Dependencies
+- For My Compatible
+```bash
+pip install -r requirements/requirements_compatible.txt
+```
+- For `End2end Inference`
+```bash
+pip install -r requirements/requirements_inference.txt
+```
+### Download Model Checkpoints
+```bash
+python scripts/download_ckpts.py
+```
+### Setup Third Party
+```bash
+python scripts/setup_third_party.py
+```
+## Data Preparation
+
+## Pre-trained Model preparation
+For detailed instructions on how to process and manage your data effectively, refer to the [Full guide for data processing](./docs/data_processing.md).
+
+This is fast usage for dataset preparation
+```bash
+python scripts/process_dataset.py
+```
+## Training
+
+1. Configure the default settings for Accelerate:
+```bash
+accelerate config default
+```
+
+2. Launch the training script using Accelerate:
+```bash
+accelerate launch src/slimface/training/accelerate_train.py
+```
+
+For additional help, you can refer to the [Training Documentation](./docs/training/training_docs.md) for more details.
+
+## Demostration
+```bash
+python apps/gradio_app.py
+```
+
+https://huggingface.co/spaces/danhtran2mind/slimface
+
+## Project Description
+
+This repository is trained from [![GitHub Repo](https://img.shields.io/badge/GitHub-danhtran2mind%2Fedgeface-blue?style=flat)](https://github.com/danhtran2mind/edgeface), a fork of [![GitHub Repo](https://img.shields.io/badge/GitHub-otroshi%2Fedgeface-blue?style=flat)](https://github.com/otroshi/edgeface), with numerous bug fixes and rewritten code for improved performance and stability.

SlimFace/apps/gradio_app.py

@@ -0,0 +1,26 @@
+import gradio as gr
+from gradio_app.inference import run_inference
+
+def create_gradio_interface():
+    return gr.Interface(
+        fn=run_inference,
+        inputs=[
+            gr.Image(type="pil", label="Upload Image"),
+            gr.File(label="Reference Dict JSON File"),
+            gr.File(label="Index to Class Mapping JSON File"),
+            gr.File(label="Classifier Model (.pth) File"),
+            gr.Textbox(label="EdgeFace Model Name", value="edgeface_base"),
+            gr.Textbox(label="EdgeFace Model Directory", value="ckpts/idiap"),
+            gr.Dropdown(choices=["yolo", "mtcnn"], label="Face Detection Algorithm", value="yolo"),
+            gr.Dropdown(choices=["auto", "cpu", "gpu"], label="Accelerator", value="auto"),
+            gr.Slider(minimum=112, maximum=448, step=1, value=224, label="Resolution"),
+            gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.6, label="Similarity Threshold")
+        ],
+        outputs="text",
+        title="Face Classification with EdgeFace Validation",
+        description="Upload an image and required files to perform face classification with EdgeFace embedding validation."
+    )
+
+if __name__ == "__main__":
+    iface = create_gradio_interface()
+    iface.launch()

SlimFace/apps/gradio_app/inference.py

@@ -0,0 +1,57 @@
+import os
+import sys
+from PIL import Image
+
+# Append the path to the inference script's directory
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'src', 'slimface', 'inference')))
+from end2end_inference import cinference_and_confirm
+
+def run_inference(image, reference_dict_path, index_to_class_mapping_path, model_path, 
+                 edgeface_model_name="edgeface_base", edgeface_model_dir="ckpts/idiap", 
+                 algorithm="yolo", accelerator="auto", resolution=224, similarity_threshold=0.6):
+    # Save uploaded image temporarily in apps/gradio_app/
+    temp_image_path = os.path.join(os.path.dirname(__file__), "temp_image.jpg")
+    image.save(temp_image_path)
+
+    # Create args object to mimic command-line arguments
+    class Args:
+        def __init__(self):
+            self.unknown_image_path = temp_image_path
+            self.reference_dict_path = reference_dict_path.name if reference_dict_path else None
+            self.index_to_class_mapping_path = index_to_class_mapping_path.name if index_to_class_mapping_path else None
+            self.model_path = model_path.name if model_path else None
+            self.edgeface_model_name = edgeface_model_name
+            self.edgeface_model_dir = edgeface_model_dir
+            self.algorithm = algorithm
+            self.accelerator = accelerator
+            self.resolution = resolution
+            self.similarity_threshold = similarity_threshold
+
+    args = Args()
+
+    # Validate inputs
+    if not all([args.reference_dict_path, args.index_to_class_mapping_path, args.model_path]):
+        return "Error: Please provide all required files (reference dict, index-to-class mapping, and model)."
+
+    try:
+        # Call the inference function from end2end_inference.py
+        results = cinference_and_confirm(args)
+        
+        # Format output
+        output = ""
+        for result in results:
+            output += f"Image: {result['image_path']}\n"
+            output += f"Predicted Class: {result['predicted_class']}\n"
+            output += f"Confidence: {result['confidence']:.4f}\n"
+            output += f"Similarity: {result.get('similarity', 'N/A'):.4f}\n"
+            output += f"Confirmed: {result.get('confirmed', 'N/A')}\n\n"
+        
+        return output
+    
+    except Exception as e:
+        return f"Error: {str(e)}"
+    
+    finally:
+        # Clean up temporary image
+        if os.path.exists(temp_image_path):
+            os.remove(temp_image_path)

SlimFace/assets/comparision.md

@@ -0,0 +1,11 @@
+slim_face_vit_b_16
+    Train loss: 0.0074, Train acc: 0.9980, Val loss: 0.2179, Val acc: 0.9336
+    
+efficientnet_b3
+    Train loss: 0.0014, Train acc: 1.0000, Val loss: 0.1931, Val acc: 0.9427
+
+efficientnet_v2_s
+    Train loss: 0.0016, Train acc: 1.0000, Val loss: 0.2374, Val acc: 0.9375
+
+regnet_y_800mf
+    Train loss: 0.0033, Train acc: 0.9997, Val loss: 0.3766, Val acc: 0.8906

SlimFace/ckpts/.gitignore

@@ -0,0 +1,4 @@
+# Ignore everything in this directory
+*
+# Except this .gitignore file
+!.gitignore

SlimFace/configs/accelerate_config.yaml

@@ -0,0 +1,6 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: FSDP
+num_processes: 4
+mixed_precision: fp16
+fsdp_config:
+  fsdp_offload_params: true

SlimFace/configs/image_classification_models_config.yaml

@@ -0,0 +1,249 @@
+# For more details on models, see https://pytorch.org/vision/main/models.html
+
+# EfficientNet models: Designed for efficiency with compound scaling of depth, width, and resolution.
+# These models balance accuracy and computational efficiency, ideal for resource-constrained environments.
+efficientnet_b0:
+  metrics:
+    Acc@1: 77.692  # Top-1 accuracy on ImageNet
+    Acc@5: 93.532  # Top-5 accuracy on ImageNet
+    GFLOPS: 0.39   # Computational complexity
+    Params: 5.3M   # Number of parameters
+  model_fn: models.efficientnet_b0
+  resolution: 224  # Input image resolution
+  weights: models.EfficientNet_B0_Weights.IMAGENET1K_V1  # Pretrained weights on ImageNet
+
+efficientnet_b1:
+  metrics:
+    Acc@1: 78.642
+    Acc@5: 94.186
+    GFLOPS: 0.69
+    Params: 7.8M
+  model_fn: models.efficientnet_b1
+  resolution: 240
+  weights: models.EfficientNet_B1_Weights.IMAGENET1K_V1
+
+efficientnet_b2:
+  metrics:
+    Acc@1: 80.608
+    Acc@5: 95.31
+    GFLOPS: 1.09
+    Params: 9.1M
+  model_fn: models.efficientnet_b2
+  resolution: 260
+  weights: models.EfficientNet_B2_Weights.IMAGENET1K_V1
+
+efficientnet_b3:
+  metrics:
+    Acc@1: 82.008
+    Acc@5: 96.054
+    GFLOPS: 1.83
+    Params: 12.2M
+  model_fn: models.efficientnet_b3
+  resolution: 300
+  weights: models.EfficientNet_B3_Weights.IMAGENET1K_V1
+
+efficientnet_b4:
+  metrics:
+    Acc@1: 83.384
+    Acc@5: 96.594
+    GFLOPS: 4.39
+    Params: 19.3M
+  model_fn: models.efficientnet_b4
+  resolution: 380
+  weights: models.EfficientNet_B4_Weights.IMAGENET1K_V1
+
+efficientnet_b5:
+  metrics:
+    Acc@1: 83.444
+    Acc@5: 96.628
+    GFLOPS: 10.27
+    Params: 30.4M
+  model_fn: models.efficientnet_b5
+  resolution: 456
+  weights: models.EfficientNet_B5_Weights.IMAGENET1K_V1
+
+efficientnet_b6:
+  metrics:
+    Acc@1: 84.008
+    Acc@5: 96.916
+    GFLOPS: 19.07
+    Params: 43.0M
+  model_fn: models.efficientnet_b6
+  resolution: 528
+  weights: models.EfficientNet_B6_Weights.IMAGENET1K_V1
+
+efficientnet_b7:
+  metrics:
+    Acc@1: 84.122
+    Acc@5: 96.908
+    GFLOPS: 37.75
+    Params: 66.3M
+  model_fn: models.efficientnet_b7
+  resolution: 600
+  weights: models.EfficientNet_B7_Weights.IMAGENET1K_V1
+
+# EfficientNet V2 models: Improved training efficiency and performance over V1.
+# These models use progressive learning and optimized scaling for better accuracy.
+efficientnet_v2_l:
+  metrics:
+    Acc@1: 85.808
+    Acc@5: 97.788
+    GFLOPS: 56.08
+    Params: 118.5M
+  model_fn: models.efficientnet_v2_l
+  resolution: 480
+  weights: models.EfficientNet_V2_L_Weights.IMAGENET1K_V1
+
+efficientnet_v2_m:
+  metrics:
+    Acc@1: 85.112
+    Acc@5: 97.156
+    GFLOPS: 24.58
+    Params: 54.1M
+  model_fn: models.efficientnet_v2_m
+  resolution: 480
+  weights: models.EfficientNet_V2_M_Weights.IMAGENET1K_V1
+
+efficientnet_v2_s:
+  metrics:
+    Acc@1: 84.228
+    Acc@5: 96.878
+    GFLOPS: 8.37
+    Params: 21.5M
+  model_fn: models.efficientnet_v2_s
+  resolution: 384
+  weights: models.EfficientNet_V2_S_Weights.IMAGENET1K_V1
+
+# RegNet models: Designed for scalability and efficiency with a focus on network design.
+# These models optimize for both accuracy and computational efficiency.
+regnet_y_128gf:
+  metrics:
+    Acc@1: 86.068  # High accuracy but computationally expensive
+    Acc@5: 97.844
+    GFLOPS: 127.52
+    Params: 644.8M
+  model_fn: models.regnet_y_128gf
+  resolution: 224
+  weights: models.RegNet_Y_128GF_Weights.IMAGENET1K_SWAG_LINEAR_V1
+
+regnet_y_16gf:
+  metrics:
+    Acc@1: 82.886
+    Acc@5: 96.328
+    GFLOPS: 15.91
+    Params: 83.6M
+  model_fn: models.regnet_y_16gf
+  resolution: 224
+  weights: models.RegNet_Y_16GF_Weights.IMAGENET1K_V2
+
+regnet_y_1_6gf:
+  metrics:
+    Acc@1: 80.876
+    Acc@5: 95.444
+    GFLOPS: 1.61
+    Params: 11.2M
+  model_fn: models.regnet_y_1_6gf
+  resolution: 224
+  weights: models.RegNet_Y_1_6GF_Weights.IMAGENET1K_V2
+
+regnet_y_32gf:
+  metrics:
+    Acc@1: 83.368
+    Acc@5: 96.498
+    GFLOPS: 32.28
+    Params: 145.0M
+  model_fn: models.regnet_y_32gf
+  resolution: 224
+  weights: models.RegNet_Y_32GF_Weights.IMAGENET1K_V2
+
+regnet_y_3_2gf:
+  metrics:
+    Acc@1: 81.982
+    Acc@5: 95.972
+    GFLOPS: 3.18
+    Params: 19.4M
+  model_fn: models.regnet_y_3_2gf
+  resolution: 224
+  weights: models.RegNet_Y_3_2GF_Weights.IMAGENET1K_V2
+
+regnet_y_400mf:
+  metrics:
+    Acc@1: 75.804
+    Acc@5: 92.742
+    GFLOPS: 0.4
+    Params: 4.3M
+  model_fn: models.regnet_y_400mf
+  resolution: 224
+  weights: models.RegNet_Y_400MF_Weights.IMAGENET1K_V2
+
+regnet_y_800mf:
+  metrics:
+    Acc@1: 78.828
+    Acc@5: 94.502
+    GFLOPS: 0.83
+    Params: 6.4M
+  model_fn: models.regnet_y_800mf
+  resolution: 224
+  weights: models.RegNet_Y_800MF_Weights.IMAGENET1K_V2
+
+regnet_y_8gf:
+  metrics:
+    Acc@1: 82.828
+    Acc@5: 96.33
+    GFLOPS: 8.47
+    Params: 39.4M
+  model_fn: models.regnet_y_8gf
+  resolution: 224
+  weights: models.RegNet_Y_8GF_Weights.IMAGENET1K_V2
+
+# Vision Transformer (ViT) models: Transformer-based architecture for image classification.
+# These models excel in capturing long-range dependencies but require significant compute for larger variants.
+vit_b_16:
+  metrics:
+    Acc@1: 81.072  # Base ViT model with balanced accuracy and efficiency
+    Acc@5: 95.318
+    GFLOPS: 17.56
+    Params: 86.6M
+  model_fn: models.vit_b_16
+  resolution: 224
+  weights: models.ViT_B_16_Weights.IMAGENET1K_V1
+
+vit_b_32:
+  metrics:
+    Acc@1: 75.912  # Smaller patch size version of ViT, lower accuracy but fewer computations
+    Acc@5: 92.466
+    GFLOPS: 4.41
+    Params: 88.2M
+  model_fn: models.vit_b_32
+  resolution: 224
+  weights: models.ViT_B_32_Weights.IMAGENET1K_V1
+
+vit_h_14:
+  metrics:
+    Acc@1: 88.552  # High-performance ViT model with very high accuracy and computational cost
+    Acc@5: 98.694
+    GFLOPS: 1016.72
+    Params: 633.5M
+  model_fn: models.vit_h_14
+  resolution: 224
+  weights: models.ViT_H_14_Weights.IMAGENET1K_SWAG_E2E_V1
+
+vit_l_16:
+  metrics:
+    Acc@1: 79.662  # Larger ViT model with improved accuracy over base models
+    Acc@5: 94.638
+    GFLOPS: 61.55
+    Params: 304.3M
+  model_fn: models.vit_l_16
+  resolution: 224
+  weights: models.ViT_L_16_Weights.IMAGENET1K_V1
+
+vit_l_32:
+  metrics:
+    Acc@1: 76.972  # Larger ViT with larger patch size, trading accuracy for reduced compute
+    Acc@5: 93.07
+    GFLOPS: 15.38
+    Params: 306.5M
+  model_fn: models.vit_l_32
+  resolution: 224
+  weights: models.ViT_L_32_Weights.IMAGENET1K_V1

SlimFace/data/.gitignore

@@ -0,0 +1,5 @@
+# Ignore everything in this directory
+*
+# Except this .gitignore file
+!.gitignore
+!raw

SlimFace/data/raw/.gitignore

@@ -0,0 +1,4 @@
+# Ignore everything in this directory
+*
+# Except this .gitignore file
+!.gitignore

SlimFace/docs/README.md

@@ -0,0 +1 @@
+# This is Docs

SlimFace/docs/data/data_processing.md

@@ -0,0 +1,135 @@
+# Data Processing for slimface Training 🖼️
+
+## Table of Contents
+
+- [Data Processing for slimface Training 🖼️](#data-processing-for-slimface-training-)
+  - [Command-Line Arguments](#command-line-arguments)
+    - [Command-Line Arguments for `process_dataset.py`](#command-line-arguments-for-process_datasetpy)
+    - [Example Usage](#example-usage)
+  - [Step-by-step process for handling a dataset](#step-by-step-process-for-handling-a-dataset)
+    - [Step 1: Clone the Repository](#step-1-clone-the-repository)
+    - [Step 2: Process the Dataset](#step-2-process-the-dataset)
+      - [Option 1: Using Dataset from Kaggle](#option-1-using-dataset-from-kaggle)
+      - [Option 2: Using a Custom Dataset](#option-2-using-a-custom-dataset)
+
+## Command-Line Arguments
+### Command-Line Arguments for `process_dataset.py`
+When running `python scripts/process_dataset.py`, you can customize the dataset processing with the following command-line arguments:
+
+| Argument | Type | Default | Description |
+|----------|------|---------|-------------|
+| `--dataset_slug` | `str` | `vasukipatel/face-recognition-dataset` | The Kaggle dataset slug in `username/dataset-name` format. Specifies which dataset to download from Kaggle. |
+| `--base_dir` | `str` | `./data` | The base directory where the dataset will be stored and processed. |
+| `--augment` | `flag` | `False` | Enables data augmentation (e.g., flipping, rotation) for training images to increase dataset variety. Use `--augment` to enable. |
+| `--random_state` | `int` | `42` | Random seed for reproducibility in the train-test split. Ensures consistent splitting across runs. |
+| `--test_split_rate` | `float` | `0.2` | Proportion of data to use for validation (between 0 and 1). For example, `0.2` means 20% of the data is used for validation. |
+| `--rotation_range` | `int` | `15` | Maximum rotation angle in degrees for data augmentation (if `--augment` is enabled). Images may be rotated randomly within this range. |
+| `--source_subdir` | `str` | `Original Images/Original Images` | Subdirectory within `raw_dir` containing the images to process. Used for both Kaggle and custom datasets. |
+| `--delete_raw` | `flag` | `False` | Deletes the raw folder after processing to save storage. Use `--delete_raw` to enable. |
+
+### Example Usage
+To process a Kaggle dataset with augmentation and a custom validation split:
+
+```bash
+python scripts/process_dataset.py \
+    --augment \
+    --test_split_rate 0.3 \
+    --rotation_range 15
+```
+
+To process a **custom dataset** with a specific subdirectory and delete the raw folder:
+
+```bash
+python scripts/process_dataset.py \
+    --source_subdir your_custom_dataset_dir \
+    --delete_raw
+```
+## Step-by-step process for handling a dataset
+These options allow flexible dataset processing tailored to your needs. 🚀
+
+### Step 1: Clone the Repository
+Ensure the `slimface` project is set up by cloning the repository and navigating to the project directory:
+
+```bash
+git clone https://github.com/danhtran2mind/slimface/
+cd slimface
+```
+
+### Step 2: Process the Dataset
+
+#### Option 1: Using Dataset from Kaggle
+To download and process the sample dataset from Kaggle, run:
+
+```bash
+python scripts/process_dataset.py
+```
+
+This script organizes the dataset into the following structure under `data/`:
+
+```markdown
+data/
+├── processed_ds/
+│   ├── train_data/
+│   │   ├── Charlize Theron/
+│   │   │   ├── Charlize Theron_70.jpg
+│   │   │   ├── Charlize Theron_46.jpg
+│   │   │   ...
+│   │   ├── Dwayne Johnson/
+│   │   │   ├── Dwayne Johnson_58.jpg
+│   │   │   ├── Dwayne Johnson_9.jpg
+│   │   │   ...
+│   └── val_data/
+│       ├── Charlize Theron/
+│       │   ├── Charlize Theron_60.jpg
+│       │   ├── Charlize Theron_45.jpg
+│       │   ...
+│       ├── Dwayne Johnson/
+│       │   ├── Dwayne Johnson_11.jpg
+│       │   ├── Dwayne Johnson_46.jpg
+│       │   ...
+├── raw/
+│   ├── Faces/
+│   │   ├── Jessica Alba_90.jpg
+│   │   ├── Hugh Jackman_70.jpg
+│   │   ...
+│   ├── Original Images/
+│   │   ├── Charlize Theron/
+│   │   │   ├── Charlize Theron_60.jpg
+│   │   │   ├── Charlize Theron_70.jpg
+│   │   │   ...
+│   │   ├── Dwayne Johnson/
+│   │   │   ├── Dwayne Johnson_11.jpg
+│   │   │   ├── Dwayne Johnson_58.jpg
+│   │   │   ...
+│   ├── dataset.zip
+│   └── Dataset.csv
+└── .gitignore
+```
+
+#### Option 2: Using a Custom Dataset
+If you prefer to use your own dataset, place it in `./data/raw/your_custom_dataset_dir/` with the following structure:
+
+```markdown
+data/
+├── raw/
+│   ├── your_custom_dataset_dir/
+│   │   ├── Charlize Theron/
+│   │   │   ├── Charlize Theron_60.jpg
+│   │   │   ├── Charlize Theron_70.jpg
+│   │   │   ...
+│   │   ├── Dwayne Johnson/
+│   │   │   ├── Dwayne Johnson_11.jpg
+│   │   │   ├── Dwayne Johnson_58.jpg
+│   │   │   ...
+```
+
+If you use your dataset, you do not need to include only human faces, because **we support face extraction using face detection**, and all extracted faces are saved at `data/processed_ds`.
+
+Then, process your custom dataset by specifying the subdirectory:
+
+```bash
+python scripts/process_dataset.py \
+    --source_subdir your_custom_dataset_dir
+```
+
+This ensures your dataset is properly formatted for training. 🚀

SlimFace/docs/inference/inference_doc.md

@@ -0,0 +1,15 @@
+```bash
+python src/slim_face/inference/inference.py \
+    --input_path <image_path> \
+    --model_path <model_path> \
+    --index_to_class_mapping_path <index_to_class_mapping_json_path>
+```
+
+## Example Usage
+
+```bash
+python src/slim_face/inference/inference.py \
+    --input_path "assets/test_images/Elon_Musk.jpg" \
+    --model_path "ckpts/slim_face_regnet_y_800mf_full_model.pth" \
+    --index_to_class_mapping_path ckpts/index_to_class_mapping.json
+```

SlimFace/docs/test/inference_test_doc.md

@@ -0,0 +1,96 @@
+# Running Inference Test Script
+
+Instructions to run the `tests/inference_test.sh` script in the `tests` folder on Linux, Windows, and macOS.
+
+## Prerequisites
+1. Install Python, PyTorch, Accelerate:
+   ```bash
+   pip install -r requirements/requirements.txt
+   ```
+2. Create virtual environment:
+   ```bash
+   python -m venv venv
+   source venv/bin/activate  # Linux/macOS
+   .\venv\Scripts\activate   # Windows
+   ```
+3. Make script executable (Linux/macOS):
+   ```bash
+   chmod +x tests/inference_test.sh
+   ```
+
+## Linux
+1. Open terminal, go to folder:
+   ```bash
+   cd tests
+   ```
+2. Run script:
+   ```bash
+   ./inference_test.sh
+   ```
+3. **Fix issues**:
+   - Use `bash inference_test.sh` if `./` fails.
+   - Fix line endings:
+     ```bash
+     sudo apt install dos2unix
+     dos2unix inference_test.sh
+     ```
+
+## Windows (using WSL)
+1. Install WSL and Ubuntu from Microsoft Store.
+2. Install dependencies:
+   ```bash
+   sudo apt update
+   sudo apt install python3 python3-pip
+   pip install -r requirements/requirements.txt
+   ```
+3. Go to folder:
+   ```bash
+   cd ./tests
+   ```
+4. Make executable:
+   ```bash
+   chmod +x inference_test.sh
+   ```
+5. Run script:
+   ```bash
+   ./inference_test.sh
+   ```
+6. **Fix issues**:
+   - Fix line endings:
+     ```bash
+     sudo apt install dos2unix
+     dos2unix inference_test.sh
+     ```
+
+## macOS
+1. Open Terminal, go to folder:
+   ```bash
+   cd tests
+   ```
+2. Install dependencies:
+   ```bash
+   brew install python
+   pip install -r requirements/requirements.txt
+   ```
+3. Make executable:
+   ```bash
+   chmod +x inference_test.sh
+   ```
+4. Run script:
+   ```bash
+   ./inference_test.sh
+   ```
+5. **Fix issues**:
+   - Fix line endings:
+     ```bash
+     brew install dos2unix
+     dos2unix inference_test.sh
+     ```
+
+## Notes
+- Ensure GPU support (CUDA for Linux/Windows, MPS for macOS) if needed.
+- Check script for extra settings (e.g., `export CUDA_VISIBLE_DEVICES=0`).
+- Save output:
+  ```bash
+  ./inference_test.sh > output.log 2>&1
+  ```

SlimFace/docs/test/training_test_doc.md

@@ -0,0 +1,103 @@
+# Running Training Scripts
+
+Instructions to run these scripts in the `tests` folder on Linux, Windows, and macOS:
+- `tests/training_accelerate_efficientnet_b3.sh`
+- `tests/training_accelerate_efficientnet_v2_s.sh`
+- `tests/training_accelerate_regnet_y_800mf.sh`
+- `tests/training_accelerate_vit_b_16_test.sh`
+
+## Prerequisites
+1. Install Python, PyTorch, Accelerate:
+   ```bash
+   pip install requirements/requirements.txt
+   ```
+2. Create virtual environment:
+   ```bash
+   python -m venv venv
+   source venv/bin/activate  # Linux/macOS
+   .\venv\Scripts\activate   # Windows
+   ```
+3. Make scripts executable (Linux/macOS):
+   ```bash
+   chmod +x tests/*.sh
+   ```
+
+## Linux
+1. Open terminal, go to folder:
+   ```bash
+   cd tests
+   ```
+2. Run scripts:
+   ```bash
+   ./training_accelerate_efficientnet_b3.sh
+   ./training_accelerate_efficientnet_v2_s.sh
+   ./training_accelerate_regnet_y_800mf.sh
+   ./training_accelerate_vit_b_16_test.sh
+   ```
+3. **Fix issues**:
+   - Use `bash training_accelerate_efficientnet_b3.sh` if `./` fails.
+   - Fix line endings:
+     ```bash
+     sudo apt install dos2unix
+     dos2unix training_accelerate_*.sh
+     ```
+
+## Windows (using WSL)
+1. Install WSL and Ubuntu from Microsoft Store.
+2. Install dependencies:
+   ```bash
+   sudo apt update
+   sudo apt install python3 python3-pip
+   pip install -r requirements/requirements.txt
+   ```
+3. Go to folder:
+   ```bash
+   cd ./tests
+   ```
+4. Make executable:
+   ```bash
+   chmod +x training_accelerate_*.sh
+   ```
+5. Run scripts:
+   ```bash
+   ./training_accelerate_efficientnet_b3.sh
+   ```
+6. **Fix issues**:
+   - Fix line endings:
+     ```bash
+     sudo apt install dos2unix
+     dos2unix training_accelerate_*.sh
+     ```
+
+## macOS
+1. Open Terminal, go to folder:
+   ```bash
+   cd tests
+   ```
+2. Install dependencies:
+   ```bash
+   brew install python
+   pip install -r requirements/requirements.txt
+   ```
+3. Make executable:
+   ```bash
+   chmod +x training_accelerate_*.sh
+   ```
+4. Run scripts:
+   ```bash
+   ./training_accelerate_efficientnet_b3.sh
+   ```
+5. **Fix issues**:
+   - Fix line endings:
+     ```bash
+     brew install dos2unix
+     dos2unix training_accelerate_*.sh
+     ```
+
+## Notes
+- Ensure GPU support (CUDA for Linux/Windows, MPS for macOS) if needed.
+- Check scripts for extra settings (e.g., `export CUDA_VISIBLE_DEVICES=0`).
+- Save output:
+  ```bash
+  ./training_accelerate_efficientnet_b3.sh > output.log 2>&1
+  ```

SlimFace/docs/training/training_doc.md

@@ -0,0 +1,48 @@
+
+# Training Documentation
+
+This document outlines the command-line arguments and a concise overview of the training pipeline for a face classification model using PyTorch Lightning.
+
+## Table of Contents
+
+-   Arguments Table
+    
+-   Training Pipeline Overview
+    
+
+# Training Arguments Documentation
+
+This document outlines the command-line arguments and a concise overview of the training pipeline for a face classification model using PyTorch Lightning.
+
+## Table of Contents
+
+- [Arguments Table](#arguments-table)
+- [Training Pipeline Overview](#training-pipeline-overview)
+
+## Arguments Table
+
+| Argument Name                          | Type  | Description                                                                                                                    |
+|----------------------------------------|-------|-------------------------------------------------------------------------------------------------------------------------------|
+| `dataset_dir`                          | `str` | Path to the dataset directory containing `train_data` and `val_data` subdirectories with preprocessed face images organized by person. |
+| `image_classification_models_config_path` | `str` | Path to the YAML configuration file defining model configurations, including model function, resolution, and weights.              |
+| `batch_size`                           | `int` | Batch size for training and validation data loaders. Affects memory usage and training speed.                                  |
+| `num_epochs`                           | `int` | Number of epochs for training the model. An epoch is one full pass through the training dataset.                               |
+| `learning_rate`                        | `float` | Initial learning rate for the Adam optimizer used during training.                                                             |
+| `max_lr_factor`                        | `float` | Multiplies the initial learning rate to determine the maximum learning rate during the warmup phase of the scheduler.           |
+| `accelerator`                          | `str` | Type of accelerator for training. Options: `cpu`, `gpu`, `tpu`, `auto`. `auto` selects the best available device.              |
+| `devices`                              | `int` | Number of devices (e.g., GPUs) to use for training. Relevant for multi-GPU training.                                           |
+| `algorithm`                            | `str` | Face detection algorithm for preprocessing images. Options: `mtcnn`, `yolo`.                                                  |
+| `warmup_steps`                         | `float` | Fraction of total training steps for the warmup phase of the learning rate scheduler (e.g., `0.05` means 5% of total steps).  |
+| `total_steps`                          | `int` | Total number of training steps. If `0`, calculated as epochs × steps per epoch (based on dataset size and batch size).         |
+| `classification_model_name`             | `str` | Name of the classification model to use, as defined in the YAML configuration file.                                            |
+
+## Training Pipeline Overview
+
+The training pipeline preprocesses face images, fine-tunes a classification head on a pretrained model, and trains using PyTorch Lightning. Key components:
+
+1. **Preprocessing**: Aligns faces using `yolo` or `mtcnn`, caches resized images (`preprocess_and_cache_images`).
+2. **Dataset**: `FaceDataset` loads pre-aligned images, applies normalization, and assigns labels by person.
+3. **Model**: `FaceClassifier` pairs a frozen pretrained model (e.g., EfficientNet) with a custom classification head.
+4. **Training**: `FaceClassifierLightning` manages training with Adam optimizer, cosine annealing scheduler, and logs loss/accuracy.
+5. **Configuration**: Loads model details from YAML (`load_model_configs`), uses `DataLoader` with multiprocessing, and saves models via `CustomModelCheckpoint`.
+6. **Execution**: `main` orchestrates preprocessing, data loading, model training, and saves full model and classifier head.

SlimFace/new_contruct.md

@@ -0,0 +1,63 @@
+```markdown
+slim-face-recognition/
+├── src/                     # Source code for the project
+│   ├── slim_face/           # Main package for your project
+│   │   ├── __init__.py      # Marks directory as a Python package
+│   │   ├── models/          # Model definitions and architectures
+│   │   │   ├── __init__.py
+│   │   │   ├── edgeface.py  # Model definitions (e.g., edgeface backbones)
+│   │   │   └── ...          # Other model-related scripts
+│   │   ├── data/            # Data loading and preprocessing
+│   │   │   ├── __init__.py
+│   │   │   ├── dataset.py   # Custom Dataset classes for DataLoader
+│   │   │   ├── align.py     # Face alignment utilities (e.g., from edgeface)
+│   │   │   └── ...          # Other data-related scripts
+│   │   ├── training/        # Training-related scripts and logic
+│   │   │   ├── __init__.py
+│   │   │   ├── train.py     # Main training script
+│   │   │   ├── accelerate_train.py  # Accelerated training script
+│   │   │   └── ...          # Other training utilities
+│   │   ├── inference/       # Inference-related scripts and logic
+│   │   │   ├── __init__.py  # Marks directory as a Python package
+│   │   │   ├── inference.py # Face recognition inference logic
+│   │   ├── utils/           # Utility functions (e.g., logging, metrics)
+│   │   │   ├── __init__.py
+│   │   │   ├── helpers.py   # Miscellaneous helper functions
+│   │   │   └── ...          # Other utility scripts
+│   │   └── __main__.py      # Entry point for running the package as a module
+├── tests/                   # Unit and integration tests
+│   ├── __init__.py
+│   ├── test_data.py         # Tests for data loading
+│   ├── test_models.py       # Tests for model functionality
+│   ├── test_training.py     # Tests for training pipeline
+│   ├── test_inference.py    # Tests for inference pipeline
+│   ├── test_images/         # Sample images for testing (e.g., Elon_Musk.jpg)
+├── data/                    # Datasets and data-related files
+│   ├── raw/                 # Raw, unprocessed data
+│   ├── processed/           # Preprocessed data (e.g., aligned faces)
+│   └── external/            # External datasets (e.g., from Kaggle)
+├── scripts/                 # Standalone scripts for tasks like data download
+│   ├── download_dataset.py  # Script to download datasets (e.g., Kaggle)
+│   └── preprocess.py        # Data preprocessing scripts
+├── notebooks/               # Jupyter notebooks for exploration and analysis
+│   ├── ztest.ipynb          # Existing notebook for testing/exploration
+│   └── ...                  # Other exploratory notebooks
+├── ckpts/             # Model checkpoints and weights
+│   ├── edgeface_xs_gamma_06.pt  # Pretrained model weights
+│   ├── edgeface_s_gamma_05.pt   # Pretrained model weights
+│   └── ...                  # Other checkpoints
+├── configs/                 # Configuration files (e.g., YAML, JSON)
+│   ├── training.yaml        # Training hyperparameters
+│   └── model.yaml           # Model configurations
+├── docs/                    # Documentation files
+│   ├── api.md               # API documentation
+│   └── usage.md             # Usage instructions
+├── requirements.txt         # Main dependencies
+├── requirements_compatible.txt    # Development dependencies (e.g., testing, linting)
+├── README.md               # Project overview and setup instructions
+├── LICENSE                 # License file (e.g., MIT, Apache)
+├── .gitignore              # Git ignore file
+├── .python-version         # Python version specification (e.g., for pyenv)
+├── setup.py                # Setup script for packaging the project
+└── pyproject.toml          # Modern Python project configuration (optional)
+```

SlimFace/requirements/requirements.txt

@@ -0,0 +1,15 @@
+torch>=2.6.0
+torchvision>=0.21.0
+timm>=1.0.15
+mxnet>=1.9.1
+opencv-python>=4.10.0.84
+numpy>=1.26.0,<2.0.0
+pytorch-lightning>=2.5.1
+tqdm
+imgaug
+accelerate>=1.6.0
+scikit-learn
+pillow
+requests
+ultralytics
+huggingface-hub>=0.31.1

SlimFace/requirements/requirements_compatible.txt

@@ -0,0 +1,15 @@
+torch==2.6.0
+torchvision==0.21.0
+timm==1.0.15
+mxnet==1.9.1
+opencv-python==4.11.0.86
+numpy==1.26.4
+pillow==11.2.1
+pytorch-lightning==2.5.1
+accelerate==1.6.0
+imgaug==0.4.0
+scikit-learn==1.6.1
+pillow==11.2.1
+requests==2.32.4
+ultralytics==8.3.160
+huggingface-hub==0.31.1

SlimFace/requirements/requirements_inference.txt

@@ -0,0 +1,15 @@
+torch>=2.6.0
+torchvision>=0.21.0
+timm>=1.0.15
+mxnet>=1.9.1
+opencv-python>=4.10.0.84
+numpy>=1.26.0,<2.0.0
+ultralytics
+pytorch-lightning>=2.5.1
+tqdm
+imgaug
+accelerate>=1.6.0
+scikit-learn
+pillow
+requests
+huggingface-hub>=0.31.1

SlimFace/scripts/download_ckpts.py

@@ -0,0 +1,104 @@
+import os
+import argparse
+from huggingface_hub import snapshot_download
+
+# Model configurations for EdgeFace models
+model_configs = {
+    "edgeface_base": {
+        "repo": "idiap/EdgeFace-Base",
+        "filename": "edgeface_base.pt",
+        "local_dir": "ckpts/idiap"
+    },
+    "edgeface_s_gamma_05": {
+        "repo": "idiap/EdgeFace-S-GAMMA",
+        "filename": "edgeface_s_gamma_05.pt",
+        "local_dir": "ckpts/idiap"
+    },
+    "edgeface_xs_gamma_06": {
+        "repo": "idiap/EdgeFace-XS-GAMMA",
+        "filename": "edgeface_xs_gamma_06.pt",
+        "local_dir": "ckpts/idiap"
+    },
+    "edgeface_xxs": {
+        "repo": "idiap/EdgeFace-XXS",
+        "filename": "edgeface_xxs.pt",
+        "local_dir": "ckpts/idiap"
+    },
+    "SlimFace_efficientnet_b3": {
+        "repo": "danhtran2mind/SlimFace-sample-checkpoints",
+        "filename": "SlimFace_efficientnet_b3_full_model.pth",
+        "local_dir": "ckpts"
+    },
+    "SlimFace_efficientnet_v2_s": {
+        "repo": "danhtran2mind/SlimFace-sample-checkpoints",
+        "filename": "SlimFace_efficientnet_v2_s_full_model.pth",
+        "local_dir": "ckpts"
+    },
+    "SlimFace_regnet_y_800mf": {
+        "repo": "danhtran2mind/SlimFace-sample-checkpoints",
+        "filename": "SlimFace_regnet_y_800mf_full_model.pth",
+        "local_dir": "ckpts"
+    },
+    "SlimFace_vit_b_16": {
+        "repo": "danhtran2mind/SlimFace-sample-checkpoints",
+        "filename": "SlimFace_vit_b_16_full_model.pth",
+        "local_dir": "ckpts"
+    },
+    "SlimFace_mapping": {
+        "repo": "danhtran2mind/SlimFace-sample-checkpoints",
+        "filename": "index_to_class_mapping.json",
+        "local_dir": "ckpts"
+    }
+}
+
+def download_models(model_name=None):
+    """Download specified models from model_configs to their respective local directories.
+
+    Args:
+        model_name (str, optional): Specific model to download. If None, download all models.
+    """
+    # Determine files to download
+    if model_name:
+        if model_name not in model_configs:
+            raise ValueError(f"Model {model_name} not found in available models: {list(model_configs.keys())}")
+        configs_to_download = [model_configs[model_name]]
+    else:
+        configs_to_download = list(model_configs.values())
+
+    for config in configs_to_download:
+        repo_id = config["repo"]
+        filename = config["filename"]
+        local_dir = config["local_dir"]
+
+        # Ensure the local directory exists
+        os.makedirs(local_dir, exist_ok=True)
+
+        try:
+            snapshot_download(
+                repo_id=repo_id,
+                local_dir=local_dir,
+                local_dir_use_symlinks=False,
+                allow_patterns=[filename],
+                cache_dir=None,
+                revision="main"
+            )
+            print(f"Downloaded {filename} to {local_dir}")
+        except Exception as e:
+            print(f"Error downloading {filename}: {e}")
+
+def main():
+    """Parse command-line arguments and initiate model download."""
+    parser = argparse.ArgumentParser(description="Download models from Hugging Face Hub.")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default=None,
+        choices=list(model_configs.keys()),
+        help="Specific model to download. If not provided, all models are downloaded."
+    )
+    args = parser.parse_args()
+
+    download_models(args.model)
+
+if __name__ == "__main__":
+    main()

SlimFace/scripts/process_dataset.py

@@ -0,0 +1,242 @@
+import os
+import zipfile
+import requests
+import json
+from tqdm import tqdm
+from sklearn.model_selection import train_test_split
+import imgaug.augmenters as iaa
+import sys
+import argparse
+import shutil
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+from src.slimface.data.data_processing import process_image
+
+def download_and_split_kaggle_dataset(
+    dataset_slug,
+    base_dir="data",
+    augment=False,
+    random_state=42,
+    test_split_rate=0.2,
+    rotation_range=15,
+    source_subdir="Original Images/Original Images",
+    delete_raw=False
+):
+    """Download a Kaggle dataset, split it into train/validation sets, and process images for face recognition.
+
+    Skips downloading if ZIP exists and unzipping if raw folder contains files.
+    Optionally deletes the raw folder to save storage.
+
+    Args:
+        dataset_slug (str): Dataset slug in 'username/dataset-name' format.
+        base_dir (str): Base directory for storing dataset.
+        augment (bool): Whether to apply data augmentation to training images.
+        random_state (int): Random seed for reproducibility in train-test split.
+        test_split_rate (float): Proportion of data to use for validation (between 0 and 1).
+        rotation_range (int): Maximum rotation angle in degrees for augmentation.
+        source_subdir (str): Subdirectory within raw_dir containing images.
+        delete_raw (bool): Whether to delete the raw folder after processing to save storage.
+
+    Raises:
+        ValueError: If test_split_rate is not between 0 and 1 or dataset_slug is invalid.
+        FileNotFoundError: If source directory is not found.
+        Exception: If dataset download fails or other errors occur.
+    """
+    try:
+        # Validate test_split_rate
+        if not 0 < test_split_rate < 1:
+            raise ValueError("test_split_rate must be between 0 and 1")
+
+        # Set up directories
+        raw_dir = os.path.join(base_dir, "raw")
+        processed_dir = os.path.join(base_dir, "processed_ds")
+        train_dir = os.path.join(processed_dir, "train_data")
+        val_dir = os.path.join(processed_dir, "val_data")
+        zip_path = os.path.join(raw_dir, "dataset.zip")
+
+        os.makedirs(raw_dir, exist_ok=True)
+        os.makedirs(processed_dir, exist_ok=True)
+
+        # Check if ZIP file already exists
+        if os.path.exists(zip_path):
+            print(f"ZIP file already exists at {zip_path}, skipping download.")
+        else:
+            # Download dataset with progress bar
+            username, dataset_name = dataset_slug.split("/")
+            if not (username and dataset_name):
+                raise ValueError("Invalid dataset slug format. Expected 'username/dataset-name'")
+        
+            dataset_url = f"https://www.kaggle.com/api/v1/datasets/download/{username}/{dataset_name}"
+            print(f"Downloading dataset {dataset_slug}...")
+            response = requests.get(dataset_url, stream=True)
+            if response.status_code != 200:
+                raise Exception(f"Failed to download dataset: {response.status_code}")
+        
+            total_size = int(response.headers.get("content-length", 0))
+            with open(zip_path, "wb") as file, tqdm(
+                desc="Downloading dataset",
+                total=total_size,
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1024,
+            ) as pbar:
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:
+                        file.write(chunk)
+                        pbar.update(len(chunk))
+        
+        # # Check if raw directory contains files, excluding the ZIP file
+        # zip_filename = os.path.basename(zip_path)
+        # if os.path.exists(raw_dir) and any(file != zip_filename for file in os.listdir(raw_dir)):
+        #     print(f"Raw directory {raw_dir} already contains files, skipping extraction.")
+        # else:
+        # Extract dataset
+        print("Extracting dataset...")
+        with zipfile.ZipFile(zip_path, "r") as zip_ref:
+            zip_ref.extractall(raw_dir)
+
+        # Define source directory
+        source_dir = os.path.join(raw_dir, source_subdir)
+        if not os.path.exists(source_dir):
+            raise FileNotFoundError(f"Source directory {source_dir} not found")
+
+        # Group files by person (subfolder names)
+        person_files = {}
+        for person in os.listdir(source_dir):
+            person_dir = os.path.join(source_dir, person)
+            if os.path.isdir(person_dir):
+                person_files[person] = [
+                    f for f in os.listdir(person_dir)
+                    if os.path.isfile(os.path.join(person_dir, f))
+                    and f.lower().endswith((".png", ".jpg", ".jpeg"))
+                ]
+
+        # Define augmentation pipeline
+        if augment:
+            aug = iaa.Sequential([
+                iaa.Fliplr(p=1.0),
+                iaa.Sometimes(
+                    0.5,
+                    iaa.Affine(rotate=(-rotation_range, rotation_range))
+                ),
+            ])
+        else:
+            aug = None
+
+        # Process and split files with progress bar
+        total_files = sum(len(images) for images in person_files.values())
+        with tqdm(total=total_files, desc="Processing and copying files", unit="file") as pbar:
+            for person, images in person_files.items():
+                # Set up directories for this person
+                train_person_dir = os.path.join(train_dir, person)
+                val_person_dir = os.path.join(val_dir, person)
+                temp_dir = os.path.join(processed_dir, "temp")
+                os.makedirs(train_person_dir, exist_ok=True)
+                os.makedirs(val_person_dir, exist_ok=True)
+                os.makedirs(temp_dir, exist_ok=True)
+
+                all_image_filenames = []
+
+                # Process images and create augmentations before splitting
+                for img in images:
+                    src_path = os.path.join(source_dir, person, img)
+                    saved_images = process_image(src_path, temp_dir, aug if augment else None)
+                    all_image_filenames.extend(saved_images)
+                    pbar.update(1)
+
+                # Split all images (original and augmented) for this person
+                train_images_filenames, val_images_filenames = train_test_split(
+                    all_image_filenames,
+                    test_size=test_split_rate,
+                    random_state=random_state,
+                )
+
+                # Move images to final train/val directories
+                for img in all_image_filenames:
+                    src = os.path.join(temp_dir, img)
+                    if not os.path.exists(src):
+                        print(f"Warning: File {src} not found, skipping.")
+                        continue
+                    if img in train_images_filenames:
+                        dst = os.path.join(train_person_dir, img)
+                    else:
+                        dst = os.path.join(val_person_dir, img)
+                    os.rename(src, dst)
+
+                # Clean up temporary directory for this person
+                shutil.rmtree(temp_dir, ignore_errors=True)
+                print(f"\nCleaned up temp directory for {person}")
+
+        # Optionally delete raw folder to save storage
+        if delete_raw:
+            print(f"Deleting raw folder {raw_dir} to save storage...")
+            shutil.rmtree(raw_dir, ignore_errors=True)
+            print(f"Raw folder {raw_dir} deleted.")
+
+        print(f"Dataset {dataset_slug} downloaded, extracted, processed, and split successfully!")
+
+    except Exception as e:
+        print(f"Error processing dataset: {e}")
+        raise
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Download and process a Kaggle dataset for face recognition.")
+    parser.add_argument(
+        "--dataset_slug",
+        type=str,
+        default="vasukipatel/face-recognition-dataset",
+        help="Kaggle dataset slug in 'username/dataset-name' format"
+    )
+    parser.add_argument(
+        "--base_dir",
+        type=str,
+        default="./data",
+        help="Base directory for storing dataset"
+    )
+    parser.add_argument(
+        "--augment",
+        action="store_true",
+        help="Enable data augmentation"
+    )
+    parser.add_argument(
+        "--random_state",
+        type=int,
+        default=42,
+        help="Random seed for train-test split reproducibility"
+    )
+    parser.add_argument(
+        "--test_split_rate",
+        type=float,
+        default=0.2,
+        help="Proportion of data for validation (between 0 and 1)"
+    )
+    parser.add_argument(
+        "--rotation_range",
+        type=int,
+        default=15,
+        help="Maximum rotation angle in degrees for augmentation"
+    )
+    parser.add_argument(
+        "--source_subdir",
+        type=str,
+        default="Original Images/Original Images",
+        help="Subdirectory within raw_dir containing images"
+    )
+    parser.add_argument(
+        "--delete_raw",
+        action="store_true",
+        help="Delete the raw folder after processing to save storage"
+    )
+
+    args = parser.parse_args()
+
+    download_and_split_kaggle_dataset(
+        dataset_slug=args.dataset_slug,
+        base_dir=args.base_dir,
+        augment=args.augment,
+        random_state=args.random_state,
+        test_split_rate=args.test_split_rate,
+        rotation_range=args.rotation_range,
+        source_subdir=args.source_subdir,
+        delete_raw=args.delete_raw
+    )

SlimFace/scripts/setup_third_party.py

@@ -0,0 +1,61 @@
+import os
+import subprocess
+import sys
+import argparse
+
+def run_command(command, error_message):
+    """Run a shell command and handle errors."""
+    try:
+        subprocess.check_call(command, shell=True)
+    except subprocess.CalledProcessError:
+        print(f"Error: {error_message}")
+        sys.exit(1)
+
+def setup_edgeface(repo_url, third_party_dir, branch=None):
+    """Set up edgeface as a third-party dependency in the specified directory."""
+    edgeface_dir = os.path.join(third_party_dir, "edgeface")
+
+    # Create third_party directory if it doesn't exist
+    if not os.path.exists(third_party_dir):
+        os.makedirs(third_party_dir)
+        print(f"Created directory: {third_party_dir}")
+
+    # Clone edgeface if not already present
+    if not os.path.exists(edgeface_dir):
+        print(f"Cloning edgeface into {edgeface_dir}...")
+        clone_command = f"git clone {repo_url} {edgeface_dir}"
+        if branch:
+            clone_command = f"git clone -b {branch} {repo_url} {edgeface_dir}"
+        run_command(
+            clone_command,
+            f"Failed to clone edgeface from {repo_url}"
+        )
+    else:
+        print(f"edgeface already exists at {edgeface_dir}")
+
+    # Verify edgeface directory contains expected files
+    if os.path.exists(edgeface_dir) and os.listdir(edgeface_dir):
+        print(f"edgeface setup completed successfully at {edgeface_dir}")
+    else:
+        print(f"Error: edgeface directory is empty or invalid")
+        sys.exit(1)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Set up edgeface as a third-party dependency.")
+    parser.add_argument(
+        "--repo-url",
+        default="https://github.com/danhtran2mind/edgeface.git",
+        help="Git repository URL for edgeface (default: %(default)s)"
+    )
+    parser.add_argument(
+        "--third-party-dir",
+        default=os.path.join("src", "third_party"),
+        help="Directory to store third-party dependencies (default: %(default)s)"
+    )
+    parser.add_argument(
+        "--branch",
+        help="Git branch to clone (optional)"
+    )
+    args = parser.parse_args()
+
+    setup_edgeface(args.repo_url, args.third_party_dir, args.branch)

SlimFace/src/slimface/data/data_processing.py

@@ -0,0 +1,67 @@
+from PIL import Image
+import numpy as np
+import os
+import imgaug.augmenters as iaa
+import random
+import uuid
+
+RANDOM_RATIO = 0.5 # 0.5
+# TARGET_SIZE = (224, 224)  # Standard size for face recognition models
+
+def process_image(src_path, dest_dir, aug=None):
+    """
+    Process an image by resizing, normalizing, and optionally augmenting it.
+    Saves both raw and augmented versions of the image.
+
+    Args:
+        src_path (str): Path to the source image
+        dest_dir (str): Destination directory for the raw and augmented images
+        aug (iaa.Sequential, optional): Augmentation pipeline
+    Returns:
+        list: List of saved image filenames (raw and optionally augmented)
+    """
+    saved_images = []
+    try:
+        # Open and process image
+        img = Image.open(src_path).convert('RGB')
+        
+        # Resize image
+        # img = img.resize(TARGET_SIZE, Image.Resampling.LANCZOS)
+        
+        # Convert to numpy array and normalize
+        img_array = np.array(img) / 255.0
+        
+        # Save raw processed image
+        raw_filename = os.path.basename(src_path)
+        base, ext = os.path.splitext(raw_filename)
+        raw_dest_path = os.path.join(dest_dir, raw_filename)
+        counter = 1
+        while os.path.exists(raw_dest_path):
+            raw_filename = f"{base}_{counter}{ext}"
+            raw_dest_path = os.path.join(dest_dir, raw_filename)
+            counter += 1
+        raw_img = Image.fromarray((img_array * 255).astype(np.uint8))
+        raw_img.save(raw_dest_path, quality=100)
+        saved_images.append(raw_filename)
+        
+        # Apply augmentation if specified and save augmented image
+        if aug and random.random() <= RANDOM_RATIO:
+            img_array_aug = aug.augment_image(img_array)
+            # Clip values to ensure valid range after augmentation
+            img_array_aug = np.clip(img_array_aug, 0, 1)
+            # Convert back to image
+            aug_img = Image.fromarray((img_array_aug * 255).astype(np.uint8))
+            # Save augmented image with unique suffix
+            aug_filename = f"aug_{base}_{uuid.uuid4().hex[:8]}{ext}"
+            aug_dest_path = os.path.join(dest_dir, aug_filename)
+            aug_img.save(aug_dest_path, quality=100)
+            saved_images.append(aug_filename)
+
+    except Image.UnidentifiedImageError:
+        print(f"Error: Cannot identify image file {src_path}")
+    except OSError as e:
+        print(f"Error processing image {src_path}: {e}")
+    except Exception as e:
+        print(f"Unexpected error processing image {src_path}: {e}")
+    
+    return saved_images

SlimFace/src/slimface/data/process_face.py

@@ -0,0 +1,64 @@
+import os
+import sys
+from PIL import Image
+from tqdm import tqdm
+import warnings
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from models.detection_models import align  # Assuming this is available in your project
+
+def extract_and_save_faces(input_dir, output_dir, algorithm='yolo', resolution=224):
+    """Preprocess images using face alignment and cache them with specified resolution."""
+    if align is None:
+        raise ImportError("face_alignment package is required for preprocessing.")
+    os.makedirs(output_dir, exist_ok=True)
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=FutureWarning, message=".*rcond.*")
+        for person in sorted(os.listdir(input_dir)):
+            person_path = os.path.join(input_dir, person)
+            if not os.path.isdir(person_path):
+                continue
+            output_person_path = os.path.join(output_dir, person)
+            os.makedirs(output_person_path, exist_ok=True)
+            skipped_count = 0
+            for img_name in tqdm(os.listdir(person_path), desc=f"Processing {person}"):
+                if not img_name.endswith(('.jpg', '.jpeg', '.png')):
+                    continue
+                img_path = os.path.join(person_path, img_name)
+                output_img_path = os.path.join(output_person_path, img_name)
+                if os.path.exists(output_img_path):
+                    skipped_count += 1
+                    continue
+                try:
+                    aligned_result = align.get_aligned_face([img_path], algorithm=algorithm)
+                    aligned_image = aligned_result[0][1] if aligned_result and len(aligned_result) > 0 else None
+                    if aligned_image is None:
+                        print(f"Face detection failed for {img_path}, using resized original image")
+                        aligned_image = Image.open(img_path).convert('RGB')
+                    aligned_image = aligned_image.resize((resolution, resolution), Image.Resampling.LANCZOS)
+                    aligned_image.save(output_img_path, quality=100)
+                except Exception as e:
+                    print(f"Error processing {img_path}: {e}")
+                    aligned_image = Image.open(img_path).convert('RGB')
+                    aligned_image = aligned_image.resize((resolution, resolution), Image.Resampling.LANCZOS)
+                    aligned_image.save(output_img_path, quality=100)
+            if skipped_count > 0:
+                print(f"Skipped {skipped_count} images for {person} that were already processed.")
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(description="Preprocess and cache images with face alignment.")
+    parser.add_argument('--input_dir', type=str, required=True, help='Directory containing raw images.')
+    parser.add_argument('--output_dir', type=str, required=True, help='Directory to save preprocessed images.')
+    parser.add_argument('--algorithm', type=str, default='yolo', choices=['yolo', 'dlib'], help='Face detection algorithm to use.')
+    parser.add_argument('--resolution', type=int, default=224, help='Resolution for the output images.')
+    
+    args = parser.parse_args()
+    extract_and_save_faces(args.input_dir, args.output_dir, args.algorithm, args.resolution)
+
+    # python src/slimface/data/preprocess.py \
+    #     --input_dir "data/raw/Original Images/Original Images" \
+    #     --output_dir "data/processed/Aligned Images" \
+    #     --algorithm "yolo" \
+    #     --resolution 224

SlimFace/src/slimface/inference/end2end_inference.py

@@ -0,0 +1,143 @@
+import os
+import sys
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+import argparse
+import warnings
+import json
+
+# Append necessary paths
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "third_party")))
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from edgeface.face_alignment import align as edgeface_align
+from edgeface.backbones import get_model
+from models.detection_models import align as align_classifier
+
+def preprocess_image(image_path, algorithm='yolo', resolution=224):
+    try:
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=FutureWarning, message=".*rcond.*")
+            aligned_result = align_classifier.get_aligned_face([image_path], algorithm=algorithm)
+            aligned_image = aligned_result[0][1] if aligned_result and len(aligned_result) > 0 else Image.open(image_path).convert('RGB')
+            aligned_image = aligned_image.resize((resolution, resolution), Image.Resampling.LANCZOS)
+    except Exception as e:
+        print(f"Error processing {image_path}: {e}")
+        aligned_image = Image.open(image_path).convert('RGB').resize((resolution, resolution), Image.Resampling.LANCZOS)
+    
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    return transform(aligned_image).unsqueeze(0)
+
+def load_model(model_path):
+    try:
+        model = torch.jit.load(model_path, map_location=torch.device('cpu'))
+        model.eval()
+        return model
+    except Exception as e:
+        raise RuntimeError(f"Failed to load model from {model_path}: {e}")
+
+def load_class_mapping(index_to_class_mapping_path):
+    try:
+        with open(index_to_class_mapping_path, 'r') as f:
+            idx_to_class = json.load(f)
+        return {int(k): v for k, v in idx_to_class.items()}
+    except Exception as e:
+        raise ValueError(f"Error loading class mapping: {e}")
+
+def get_edgeface_embeddings(image_path, model_name="edgeface_base", model_dir="ckpts/idiap"):
+    model = get_model(model_name)
+    model.load_state_dict(torch.load(f'{model_dir}/{model_name}.pt', map_location='cpu'))
+    model.eval()
+    
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+    ])
+    
+    aligned_result = edgeface_align.get_aligned_face(image_path, algorithm='yolo')
+    if not aligned_result:
+        raise ValueError(f"Face alignment failed for {image_path}")
+    
+    with torch.no_grad():
+        return model(transform(aligned_result[0][1]).unsqueeze(0))
+
+def inference_and_confirm(args):
+    idx_to_class = load_class_mapping(args.index_to_class_mapping_path)
+    classifier_model = load_model(args.model_path)
+    device = torch.device('cuda' if torch.cuda.is_available() and args.accelerator == 'gpu' else 'cpu')
+    classifier_model = classifier_model.to(device)
+    
+    # Load reference images mapping from JSON file
+    try:
+        with open(args.reference_dict_path, 'r') as f:
+            reference_images = json.load(f)
+    except Exception as e:
+        raise ValueError(f"Error loading reference images from {args.reference_dict_path}: {e}")
+    
+    # Handle single image or directory
+    image_paths = [args.unknown_image_path] if args.unknown_image_path.endswith(('.jpg', '.jpeg', '.png')) else [
+        os.path.join(args.unknown_image_path, img) for img in os.listdir(args.unknown_image_path) 
+        if img.endswith(('.jpg', '.jpeg', '.png'))
+    ]
+    
+    results = []
+    with torch.no_grad():
+        for image_path in image_paths:
+            image_tensor = preprocess_image(image_path, args.algorithm, args.resolution).to(device)
+            output = classifier_model(image_tensor)
+            probabilities = torch.softmax(output, dim=1)
+            confidence, predicted = torch.max(probabilities, 1)
+            predicted_class = idx_to_class.get(predicted.item(), "Unknown")
+            
+            result = {'image_path': image_path, 'predicted_class': predicted_class, 'confidence': confidence.item()}
+            
+            # Validate with EdgeFace embeddings if reference image exists
+            reference_image_path = reference_images.get(predicted_class)
+            if reference_image_path and os.path.exists(reference_image_path):
+                unknown_embedding = get_edgeface_embeddings(image_path, args.edgeface_model_name, args.edgeface_model_dir)
+                reference_embedding = get_edgeface_embeddings(reference_image_path, args.edgeface_model_name, args.edgeface_model_dir)
+                similarity = torch.nn.functional.cosine_similarity(unknown_embedding, reference_embedding).item()
+                result['similarity'] = similarity
+                result['confirmed'] = similarity >= args.similarity_threshold
+            
+            results.append(result)
+
+    #  {'image_path': 'tests/test_images/dont_know.jpg', 'predicted_class': 'Robert Downey Jr',
+    #  'confidence': 0.9292604923248291, 'similarity': 0.603316068649292, 'confirmed': True}
+
+    return results
+
+def main(args):
+    results = inference_and_confirm(args)
+    for result in results:
+        print(f"Image: {result['image_path']}, Predicted Class: {result['predicted_class']}, "
+              f"Confidence: {result['confidence']:.4f}, Similarity: {result.get('similarity', 'N/A'):.4f}, "
+              f"Confirmed: {result.get('confirmed', 'N/A')}")
+        
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Face classification with EdgeFace embedding validation.')
+    parser.add_argument('--unknown_image_path', type=str, required=True, help='Path to image or directory.')
+    parser.add_argument('--reference_dict_path', type=str, required=True, help='Path to JSON file mapping classes to reference image paths.')
+    parser.add_argument('--index_to_class_mapping_path', type=str, required=True, help='Path to index-to-class JSON.')
+    parser.add_argument('--model_path', type=str, required=True, help='Path to classifier model (.pth).')
+    parser.add_argument('--edgeface_model_name', type=str, default='edgeface_base', help='EdgeFace model name.')
+    parser.add_argument('--edgeface_model_dir', type=str, default='ckpts/idiap', help='EdgeFace model directory.')
+    parser.add_argument('--algorithm', type=str, default='yolo', choices=['mtcnn', 'yolo'], help='Face detection algorithm.')
+    parser.add_argument('--accelerator', type=str, default='auto', choices=['cpu', 'gpu', 'auto'], help='Accelerator type.')
+    parser.add_argument('--resolution', type=int, default=224, help='Input image resolution.')
+    parser.add_argument('--similarity_threshold', type=float, default=0.6, help='Cosine similarity threshold.')
+    
+    args = parser.parse_args()
+    main(args)
+
+    # python src/slimface/inference/end2end_inference.py \
+    # --unknown_image_path tests/test_images/dont_know.jpg \
+    # --reference_dict_path tests/reference_image_data.json \
+    # --index_to_class_mapping_path /content/SlimFace/ckpts/index_to_class_mapping.json \
+    # --model_path /content/SlimFace/ckpts/SlimFace_efficientnet_b3_full_model.pth \
+    # --edgeface_model_name edgeface_base \
+    # --similarity_threshold 0.6

SlimFace/src/slimface/inference/inference.py

@@ -0,0 +1,126 @@
+import os
+import sys
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+import argparse
+import warnings
+import json
+
+# Append the parent directory's 'models/edgeface' folder to the system path
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from models.detection_models import align
+
+def preprocess_image(image_path, algorithm='yolo', resolution=224):
+    """Preprocess a single image using face alignment and specified resolution."""
+    if align is None:
+        raise ImportError("face_alignment package is required for preprocessing.")
+    try:
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=FutureWarning, message=".*rcond.*")
+            aligned_result = align.get_aligned_face([image_path], algorithm=algorithm)
+            aligned_image = aligned_result[0][1] if aligned_result and len(aligned_result) > 0 else None
+            if aligned_image is None:
+                print(f"Face detection failed for {image_path}, using resized original image")
+                aligned_image = Image.open(image_path).convert('RGB')
+            aligned_image = aligned_image.resize((resolution, resolution), Image.Resampling.LANCZOS)
+    except Exception as e:
+        print(f"Error processing {image_path}: {e}")
+        aligned_image = Image.open(image_path).convert('RGB')
+        aligned_image = aligned_image.resize((resolution, resolution), Image.Resampling.LANCZOS)
+    
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    image_tensor = transform(aligned_image).unsqueeze(0)  # Add batch dimension
+    return image_tensor
+
+def load_model(model_path):
+    """Load the trained model in TorchScript format."""
+    try:
+        model = torch.jit.load(model_path, map_location=torch.device('cpu'))
+        model.eval()
+        return model
+    except Exception as e:
+        raise RuntimeError(f"Failed to load TorchScript model from {model_path}: {e}")
+
+def load_class_mapping(index_to_class_mapping_path):
+    """Load class-to-index mapping from the JSON file."""
+    try:
+        with open(index_to_class_mapping_path, 'r') as f:
+            idx_to_class = json.load(f)
+        # Convert string keys (from JSON) to integers
+        idx_to_class = {int(k): v for k, v in idx_to_class.items()}
+        return idx_to_class
+    except FileNotFoundError:
+        raise FileNotFoundError(f"Index to class mapping file {index_to_class_mapping_path} not found.")
+    except Exception as e:
+        raise ValueError(f"Error loading index to class mapping: {e}")
+
+def inference(args):
+    # Load class mapping from JSON file
+    idx_to_class = load_class_mapping(args.index_to_class_mapping_path)
+    
+    # Load model
+    model = load_model(args.model_path)
+    
+    # Process input images
+    device = torch.device('cuda' if torch.cuda.is_available() and args.accelerator == 'gpu' else 'cpu')
+    model = model.to(device)
+    
+    image_paths = []
+    if os.path.isdir(args.input_path):
+        for img_name in os.listdir(args.input_path):
+            if img_name.endswith(('.jpg', '.jpeg', '.png')):
+                image_paths.append(os.path.join(args.input_path, img_name))
+    else:
+        if args.input_path.endswith(('.jpg', '.jpeg', '.png')):
+            image_paths.append(args.input_path)
+        else:
+            raise ValueError("Input path must be a directory or a valid image file.")
+    
+    # Perform inference
+    results = []
+    with torch.no_grad():
+        for image_path in image_paths:
+            image_tensor = preprocess_image(image_path, algorithm=args.algorithm, resolution=args.resolution)
+            image_tensor = image_tensor.to(device)
+            output = model(image_tensor)
+            probabilities = torch.softmax(output, dim=1)
+            confidence, predicted = torch.max(probabilities, 1)
+            predicted_class = idx_to_class.get(predicted.item(), "Unknown")
+            results.append({
+                'image_path': image_path,
+                'predicted_class': predicted_class,
+                'confidence': confidence.item()
+            })
+def main(args):
+    results = inference(args)
+    # Output results
+    for result in results:
+        print(f"Image: {result['image_path']}")
+        print(f"Predicted Class: {result['predicted_class']}")
+        print(f"Confidence: {result['confidence']:.4f}")
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Perform inference with a trained face classification model.')
+    parser.add_argument('--input_path', type=str, required=True,
+                        help='Path to an image or directory of images for inference.')
+    parser.add_argument('--index_to_class_mapping_path', type=str, required=True,
+                        help='Path to the JSON file containing index to class mapping.')
+    parser.add_argument('--model_path', type=str, required=True,
+                        help='Path to the trained full model in TorchScript format (.pth file).')
+    parser.add_argument('--algorithm', type=str, default='yolo',
+                        choices=['mtcnn', 'yolo'],
+                        help='Face detection algorithm to use (mtcnn or yolo).')
+    parser.add_argument('--accelerator', type=str, default='auto',
+                        choices=['cpu', 'gpu', 'auto'],
+                        help='Accelerator type for inference.')
+    parser.add_argument('--resolution', type=int, default=224,
+                        help='Resolution for input images (default: 224).')
+    
+    args = parser.parse_args()
+    main(args)

SlimFace/src/slimface/models/classification_models/alls.py

@@ -0,0 +1,55 @@
+import torch
+import torch.nn as nn
+
+class FaceClassifier(nn.Module):
+    """Face classification model with a configurable head."""
+    def __init__(self, base_model, num_classes, model_name, model_configs):
+        super(FaceClassifier, self).__init__()
+        self.base_model = base_model
+        self.model_name = model_name
+        
+        # Determine the feature extraction method and output shape
+        with torch.no_grad():
+            dummy_input = torch.zeros(1, 3, model_configs[model_name]['resolution'], model_configs[model_name]['resolution'])
+            features = base_model(dummy_input)
+            if len(features.shape) == 4:  # Spatial feature map (batch, channels, height, width)
+                in_channels = features.shape[1]
+                self.feature_type = 'spatial'
+                self.feature_dim = in_channels
+            elif len(features.shape) == 2:  # Flattened feature vector (batch, features)
+                in_channels = features.shape[1]
+                self.feature_type = 'flat'
+                self.feature_dim = in_channels
+            else:
+                raise ValueError(f"Unexpected feature shape from base model {model_name}: {features.shape}")
+
+        # Define the classifier head based on feature type
+        if self.feature_type == 'flat' or 'vit' in model_name:
+            self.conv_head = nn.Sequential(
+                nn.Linear(self.feature_dim, 512),
+                nn.BatchNorm1d(512),
+                nn.ReLU(),
+                nn.Dropout(0.5),
+                nn.Linear(512, 256),
+                nn.BatchNorm1d(256),
+                nn.ReLU(),
+                nn.Linear(256, num_classes)
+            )
+        else:
+            self.conv_head = nn.Sequential(
+                nn.Conv2d(self.feature_dim, 512, kernel_size=3, padding=1),
+                nn.BatchNorm2d(512),
+                nn.ReLU(),
+                nn.Dropout2d(0.5),
+                nn.Conv2d(512, 256, kernel_size=3, padding=1),
+                nn.BatchNorm2d(256),
+                nn.ReLU(),
+                nn.AdaptiveAvgPool2d(1),
+                nn.Flatten(),
+                nn.Linear(256, num_classes)
+            )
+
+    def forward(self, x):
+        features = self.base_model(x)
+        output = self.conv_head(features)
+        return output

SlimFace/src/slimface/models/detection_models/align.py

@@ -0,0 +1,57 @@
+import torch
+from PIL import Image
+from typing import Union, List, Tuple
+from . import mtcnn
+from .face_yolo import face_yolo_detection
+
+# Device configuration
+DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+# Initialize MTCNN model
+MTCNN_MODEL = mtcnn.MTCNN(device=DEVICE, crop_size=(112, 112))
+
+def add_image_padding(pil_img: Image.Image, top: int, right: int, bottom: int, left: int, 
+                     color: Tuple[int, int, int] = (0, 0, 0)) -> Image.Image:
+    """Add padding to a PIL image."""
+    width, height = pil_img.size
+    new_width, new_height = width + right + left, height + top + bottom
+    padded_img = Image.new(pil_img.mode, (new_width, new_height), color)
+    padded_img.paste(pil_img, (left, top))
+    return padded_img
+
+def detect_faces_mtcnn(image: Union[str, Image.Image]) -> Tuple[Union[list, None], Union[Image.Image, None]]:
+    """Detect and align faces using MTCNN model."""
+    if isinstance(image, str):
+        image = Image.open(image).convert('RGB')
+    
+    if not isinstance(image, Image.Image):
+        raise TypeError("Input must be a PIL Image or path to an image")
+    
+    try:
+        bboxes, faces = MTCNN_MODEL.align_multi(image, limit=1)
+        return bboxes[0] if bboxes else None, faces[0] if faces else None
+    except Exception as e:
+        print(f"MTCNN face detection failed: {e}")
+        return None, None
+
+def get_aligned_face(image_input: Union[str, List[str]], 
+                    algorithm: str = 'mtcnn') -> List[Tuple[Union[list, None], Union[Image.Image, None]]]:
+    """Get aligned faces from image(s) using specified algorithm."""
+    if algorithm not in ['mtcnn', 'yolo']:
+        raise ValueError("Algorithm must be 'mtcnn' or 'yolo'")
+
+    # Convert single image path to list for consistent processing
+    image_paths = [image_input] if isinstance(image_input, str) else image_input
+    if not isinstance(image_paths, list):
+        raise TypeError("Input must be a string or list of strings")
+
+    if algorithm == 'mtcnn':
+        return [detect_faces_mtcnn(path) for path in image_paths]
+    
+    # YOLO detection
+    results = face_yolo_detection(
+        image_paths,
+        use_batch=True,
+        device=DEVICE
+    )
+    return list(results)

SlimFace/src/slimface/models/detection_models/face_yolo.py

@@ -0,0 +1,151 @@
+from ultralytics import YOLO
+import cv2
+import os
+from PIL import Image
+import numpy as np
+import glob
+import sys
+import argparse
+import torch
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
+
+from utils import download_yolo_face_detection
+
+def initialize_yolo_model(yolo_model_path):
+    """Initialize YOLO model with specified device."""
+    # if device.startswith('cuda') and not torch.cuda.is_available():
+    #     print("Warning: CUDA not available, falling back to CPU.")
+    #     device = 'cpu'
+    if not os.path.exists(yolo_model_path):
+        download_yolo_face_detection.download_yolo_face_detection_model()
+    return YOLO(yolo_model_path)
+
+def process_image_results(image, image_rgb, boxes):
+    """Process bounding boxes and crop faces for a single image."""
+    bounding_boxes, cropped_faces = [], []
+    for box in boxes:
+        x1, y1, x2, y2 = map(int, box)
+        if x2 > x1 and y2 > y1 and x1 >= 0 and y1 >= 0 and x2 <= image.shape[1] and y2 <= image.shape[0]:
+            bounding_boxes.append([x1, y1, x2, y2])
+            cropped_face = image_rgb[y1:y2, x1:x2]
+            if cropped_face.size > 0:
+                pil_image = Image.fromarray(cropped_face).resize((112, 112), Image.Resampling.BILINEAR)
+                cropped_faces.append(pil_image)
+    return np.array(bounding_boxes, dtype=np.int32) if bounding_boxes else np.empty((0, 4), dtype=np.int32), cropped_faces
+
+def process_batch(model, image_paths, all_bounding_boxes, all_cropped_faces, device):
+    """Process images in batch mode using list comprehensions for efficiency."""
+    # Validate and load images, filter out invalid ones
+    valid_data = [(cv2.imread(path), path) for path in image_paths if os.path.exists(path)]
+    valid_images, valid_image_paths = zip(*[(img, path) for img, path in valid_data if img is not None]) if valid_data else ([], [])
+
+    # Append empty results for invalid images
+    for path in image_paths:
+        if not os.path.exists(path) or cv2.imread(path) is None:
+            all_bounding_boxes.append(np.empty((0, 4), dtype=np.int32))
+            all_cropped_faces.append([])
+            print(f"Warning: {'not found' if not os.path.exists(path) else 'failed to load'} {path}. Skipping.")
+
+    # Process valid images
+    if valid_images:
+        images_rgb = [cv2.cvtColor(img, cv2.COLOR_BGR2RGB) for img in valid_images]
+        results = model.predict(source=valid_image_paths, conf=0.25, iou=0.45, verbose=False, device=device)
+
+        # Process results with comprehension
+        for img, rgb, result in zip(valid_images, images_rgb, results):
+            bboxes, faces = process_image_results(img, rgb, result.boxes.xyxy.cpu().numpy())
+            all_bounding_boxes.append(bboxes)
+            all_cropped_faces.append(faces[0] if faces else [])
+
+def process_individual(model, image_paths, all_bounding_boxes, all_cropped_faces, device):
+    """Process images individually."""
+    for image_path in image_paths:
+        if not os.path.exists(image_path):
+            print(f"Warning: {image_path} not found. Skipping.")
+            all_bounding_boxes.append(np.empty((0, 4), dtype=np.int32))
+            all_cropped_faces.append([])
+            continue
+        
+        image = cv2.imread(image_path)
+        if image is None:
+            print(f"Warning: Failed to load {image_path}. Skipping.")
+            all_bounding_boxes.append(np.empty((0, 4), dtype=np.int32))
+            all_cropped_faces.append([])
+            continue
+        
+        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        results = model(image_path, conf=0.25, iou=0.45, verbose=False, device=device)
+        
+        for result in results:
+            boxes = result.boxes.xyxy.cpu().numpy()
+            bboxes, faces = process_image_results(image, image_rgb, boxes)
+            all_bounding_boxes.append(bboxes)
+            all_cropped_faces.append(faces[0] if faces else [])
+
+def face_yolo_detection(image_paths,
+                        yolo_model_path="./ckpts/yolo_face_detection/model.pt",
+                        use_batch=True, device='cuda'):
+    """Perform face detection using YOLOv11 with batch or individual processing on specified device."""
+    model = initialize_yolo_model(yolo_model_path)
+    all_bounding_boxes, all_cropped_faces = [], []
+    
+    if use_batch:
+        process_batch(model, image_paths, all_bounding_boxes, all_cropped_faces, device)
+    else:
+        process_individual(model, image_paths, all_bounding_boxes, all_cropped_faces, device)
+    
+    return zip(all_bounding_boxes, all_cropped_faces)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="YOLOv11 face detection")
+    parser.add_argument("--use-batch", action="store_true", default=True, help="Use batch processing (default: True)")
+    parser.add_argument("--image-dir", type=str, default="test/test_images", help="Input image directory")
+    parser.add_argument("--yolo-model-path", type=str, default="checkpoints/yolo11_face_detection/model.pt", help="YOLO model path")
+    parser.add_argument("--device", type=str, default="cuda", help="Device to run the model (e.g., 'cuda', 'cpu', 'cuda:0')")
+    
+    args = parser.parse_args()
+    
+    image_paths = (glob.glob(os.path.join(args.image_dir, "*.[jJ][pP][gG]")) + 
+                   glob.glob(os.path.join(args.image_dir, "*.[pP][nN][gG]")))
+    
+    if args.yolo_model_path:
+        yolo_model_path = args.yolo_model_path
+    else:
+        yolo_model_path = os.path.join("checkpoints", "yolo11_face_detection", "model.pt")
+
+    import time
+    t1 = time.time()
+    results = face_yolo_detection(image_paths, yolo_model_path, args.use_batch, args.device)
+    print("Time taken:", time.time() - t1)
+
+    # Optional: Save or process results
+    # for i, (bboxes, faces) in enumerate(results):
+    #     print(f"Image {i}: Bounding Boxes: {bboxes}")
+    #     for j, face in enumerate(faces):
+    #         face.save(f"face_{i}_{j}.png")
+
+    # Benchmarking (uncomment to use)
+    # import time
+    # num_runs = 50
+    # batch_times, individual_times = [], []
+    
+    # # Benchmark batch processing
+    # for _ in range(num_runs):
+    #     t1 = time.time()
+    #     face_yolo_detection(image_paths, yolo_model_path, use_batch=True, device=args.device)
+    #     batch_times.append(time.time() - t1)
+    
+    # # Benchmark individual processing
+    # for _ in range(num_runs):
+    #     t1 = time.time()
+    #     face_yolo_detection(image_paths, yolo_model_path, use_batch=False, device=args.device)
+    #     individual_times.append(time.time() - t1)
+    
+    # # Calculate and print average times
+    # avg_batch_time = sum(batch_times) / num_runs
+    # avg_individual_time = sum(individual_times) / num_runs
+    
+    # print(f"\nBenchmark Results (over {num_runs} runs):")
+    # print(f"Average Batch Processing Time: {avg_batch_time:.4f} seconds")
+    # print(f"Average Individual Processing Time: {avg_individual_time:.4f} seconds")

SlimFace/src/slimface/models/detection_models/mtcnn.py

@@ -0,0 +1,175 @@
+from typing import Tuple
+import numpy as np
+import torch
+from PIL import Image
+from torch.autograd import Variable
+
+import sys
+import os
+
+sys.path.insert(0, os.path.dirname(__file__))
+
+from mtcnn_pytorch.src.get_nets import PNet, RNet, ONet
+from mtcnn_pytorch.src.box_utils import nms, calibrate_box, get_image_boxes, convert_to_square
+from mtcnn_pytorch.src.first_stage import run_first_stage
+from mtcnn_pytorch.src.align_trans import get_reference_facial_points, warp_and_crop_face
+
+
+class MTCNN():
+    def __init__(self, device: str = 'cuda:0', crop_size: Tuple[int, int] = (112, 112)):
+
+        assert device in ['cuda:0', 'cpu']
+        self.device = torch.device(device)
+        assert crop_size in [(112, 112), (96, 112)]
+        self.crop_size = crop_size
+
+        # change working dir to this file location to load npz files. Then switch back
+        cwd = os.getcwd()
+        os.chdir(os.path.dirname(__file__))
+
+        self.pnet = PNet().to(self.device)
+        self.rnet = RNet().to(self.device)
+        self.onet = ONet().to(self.device)
+        self.pnet.eval()
+        self.rnet.eval()
+        self.onet.eval()
+        self.refrence = get_reference_facial_points(default_square=crop_size[0] == crop_size[1])
+
+        self.min_face_size = 20
+        self.thresholds =  [0.6,0.7,0.9]
+        self.nms_thresholds = [0.7, 0.7, 0.7]
+        self.factor = 0.85
+
+
+        os.chdir(cwd)
+
+    def align(self, img):
+        _, landmarks = self.detect_faces(img, self.min_face_size, self.thresholds, self.nms_thresholds, self.factor)
+        facial5points = [[landmarks[0][j], landmarks[0][j + 5]] for j in range(5)]
+        warped_face = warp_and_crop_face(np.array(img), facial5points, self.refrence, crop_size=self.crop_size)
+        return Image.fromarray(warped_face)
+
+    def align_multi(self, img, limit=None):
+        boxes, landmarks = self.detect_faces(img, self.min_face_size, self.thresholds, self.nms_thresholds, self.factor)
+        if limit:
+            boxes = boxes[:limit]
+            landmarks = landmarks[:limit]
+        faces = []
+        for landmark in landmarks:
+            facial5points = [[landmark[j], landmark[j + 5]] for j in range(5)]
+            warped_face = warp_and_crop_face(np.array(img), facial5points, self.refrence, crop_size=self.crop_size)
+            faces.append(Image.fromarray(warped_face))
+        return boxes, faces
+
+    def detect_faces(self, image, min_face_size, thresholds, nms_thresholds, factor):
+        """
+        Arguments:
+            image: an instance of PIL.Image.
+            min_face_size: a float number.
+            thresholds: a list of length 3.
+            nms_thresholds: a list of length 3.
+
+        Returns:
+            two float numpy arrays of shapes [n_boxes, 4] and [n_boxes, 10],
+            bounding boxes and facial landmarks.
+        """
+
+        # BUILD AN IMAGE PYRAMID
+        width, height = image.size
+        min_length = min(height, width)
+
+        min_detection_size = 12
+        # factor = 0.707  # sqrt(0.5)
+
+        # scales for scaling the image
+        scales = []
+
+        # scales the image so that
+        # minimum size that we can detect equals to
+        # minimum face size that we want to detect
+        m = min_detection_size / min_face_size
+        min_length *= m
+
+        factor_count = 0
+        while min_length > min_detection_size:
+            scales.append(m * factor**factor_count)
+            min_length *= factor
+            factor_count += 1
+
+        # STAGE 1
+
+        # it will be returned
+        bounding_boxes = []
+
+        with torch.no_grad():
+            # run P-Net on different scales
+            for s in scales:
+                boxes = run_first_stage(image, self.pnet, scale=s, threshold=thresholds[0])
+                bounding_boxes.append(boxes)
+
+            # collect boxes (and offsets, and scores) from different scales
+            bounding_boxes = [i for i in bounding_boxes if i is not None]
+            if len(bounding_boxes) == 0:
+                return [], []
+            bounding_boxes = np.vstack(bounding_boxes)
+
+            keep = nms(bounding_boxes[:, 0:5], nms_thresholds[0])
+            bounding_boxes = bounding_boxes[keep]
+
+            # use offsets predicted by pnet to transform bounding boxes
+            bounding_boxes = calibrate_box(bounding_boxes[:, 0:5], bounding_boxes[:, 5:])
+            # shape [n_boxes, 5]
+
+            bounding_boxes = convert_to_square(bounding_boxes)
+            bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
+
+            # STAGE 2
+
+            img_boxes = get_image_boxes(bounding_boxes, image, size=24)
+            img_boxes = torch.FloatTensor(img_boxes).to(self.device)
+
+            output = self.rnet(img_boxes)
+            offsets = output[0].cpu().data.numpy()  # shape [n_boxes, 4]
+            probs = output[1].cpu().data.numpy()  # shape [n_boxes, 2]
+
+            keep = np.where(probs[:, 1] > thresholds[1])[0]
+            bounding_boxes = bounding_boxes[keep]
+            bounding_boxes[:, 4] = probs[keep, 1].reshape((-1, ))
+            offsets = offsets[keep]
+
+            keep = nms(bounding_boxes, nms_thresholds[1])
+            bounding_boxes = bounding_boxes[keep]
+            bounding_boxes = calibrate_box(bounding_boxes, offsets[keep])
+            bounding_boxes = convert_to_square(bounding_boxes)
+            bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
+
+            # STAGE 3
+
+            img_boxes = get_image_boxes(bounding_boxes, image, size=48)
+            if len(img_boxes) == 0:
+                return [], []
+            img_boxes = torch.FloatTensor(img_boxes).to(self.device)
+            output = self.onet(img_boxes)
+            landmarks = output[0].cpu().data.numpy()  # shape [n_boxes, 10]
+            offsets = output[1].cpu().data.numpy()  # shape [n_boxes, 4]
+            probs = output[2].cpu().data.numpy()  # shape [n_boxes, 2]
+
+            keep = np.where(probs[:, 1] > thresholds[2])[0]
+            bounding_boxes = bounding_boxes[keep]
+            bounding_boxes[:, 4] = probs[keep, 1].reshape((-1, ))
+            offsets = offsets[keep]
+            landmarks = landmarks[keep]
+
+            # compute landmark points
+            width = bounding_boxes[:, 2] - bounding_boxes[:, 0] + 1.0
+            height = bounding_boxes[:, 3] - bounding_boxes[:, 1] + 1.0
+            xmin, ymin = bounding_boxes[:, 0], bounding_boxes[:, 1]
+            landmarks[:, 0:5] = np.expand_dims(xmin, 1) + np.expand_dims(width, 1) * landmarks[:, 0:5]
+            landmarks[:, 5:10] = np.expand_dims(ymin, 1) + np.expand_dims(height, 1) * landmarks[:, 5:10]
+
+            bounding_boxes = calibrate_box(bounding_boxes, offsets)
+            keep = nms(bounding_boxes, nms_thresholds[2], mode='min')
+            bounding_boxes = bounding_boxes[keep]
+            landmarks = landmarks[keep]
+
+        return bounding_boxes, landmarks

SlimFace/src/slimface/models/detection_models/mtcnn_pytorch/.gitignore

@@ -0,0 +1,3 @@
+.ipynb_checkpoints
+__pycache__
+

SlimFace/src/slimface/models/detection_models/mtcnn_pytorch/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2017 Dan Antoshchenko
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

SlimFace/src/slimface/models/detection_models/mtcnn_pytorch/README.md

@@ -0,0 +1,26 @@
+# MTCNN
+
+`pytorch` implementation of **inference stage** of face detection algorithm described in  
+[Joint Face Detection and Alignment using Multi-task Cascaded Convolutional Networks](https://arxiv.org/abs/1604.02878).
+
+## Example
+![example of a face detection](images/example.png)
+
+## How to use it
+Just download the repository and then do this
+```python
+from src import detect_faces
+from PIL import Image
+
+image = Image.open('image.jpg')
+bounding_boxes, landmarks = detect_faces(image)
+```
+For examples see `test_on_images.ipynb`.
+
+## Requirements
+* pytorch 0.2
+* Pillow, numpy
+
+## Credit
+This implementation is heavily inspired by:
+* [pangyupo/mxnet_mtcnn_face_detection](https://github.com/pangyupo/mxnet_mtcnn_face_detection)  

SlimFace/src/slimface/models/detection_models/mtcnn_pytorch/caffe_models/det1.prototxt

@@ -0,0 +1,177 @@
+name: "PNet"
+input: "data"
+input_dim: 1
+input_dim: 3
+input_dim: 12
+input_dim: 12
+
+layer {
+  name: "conv1"
+  type: "Convolution"
+  bottom: "data"
+  top: "conv1"
+  param {
+    lr_mult: 1
+    decay_mult: 1
+  }
+  param {
+    lr_mult: 2
+    decay_mult: 0
+  }
+  convolution_param {
+    num_output: 10
+    kernel_size: 3
+    stride: 1
+    weight_filler {
+      type: "xavier"
+    }
+    bias_filler {
+      type: "constant"
+      value: 0
+    }
+  }
+}
+layer {
+  name: "PReLU1"
+  type: "PReLU"
+  bottom: "conv1"
+  top: "conv1"
+}
+layer {
+  name: "pool1"
+  type: "Pooling"
+  bottom: "conv1"
+  top: "pool1"
+  pooling_param {
+    pool: MAX
+    kernel_size: 2
+    stride: 2
+  }
+}
+
+layer {
+  name: "conv2"
+  type: "Convolution"
+  bottom: "pool1"
+  top: "conv2"
+  param {
+    lr_mult: 1
+    decay_mult: 1
+  }
+  param {
+    lr_mult: 2
+    decay_mult: 0
+  }
+  convolution_param {
+    num_output: 16
+    kernel_size: 3
+    stride: 1
+     weight_filler {
+      type: "xavier"
+    }
+    bias_filler {
+      type: "constant"
+      value: 0
+    }
+  }
+}
+layer {
+  name: "PReLU2"
+  type: "PReLU"
+  bottom: "conv2"
+  top: "conv2"
+}
+
+layer {
+  name: "conv3"
+  type: "Convolution"
+  bottom: "conv2"
+  top: "conv3"
+  param {
+    lr_mult: 1
+    decay_mult: 1
+  }
+  param {
+    lr_mult: 2
+    decay_mult: 0
+  }
+  convolution_param {
+    num_output: 32
+    kernel_size: 3
+    stride: 1
+     weight_filler {
+      type: "xavier"
+    }
+    bias_filler {
+	  type: "constant"
+      value: 0
+    }
+  }
+}
+layer {
+  name: "PReLU3"
+  type: "PReLU"
+  bottom: "conv3"
+  top: "conv3"
+}
+
+
+layer {
+  name: "conv4-1"
+  type: "Convolution"
+  bottom: "conv3"
+  top: "conv4-1"
+  param {
+    lr_mult: 1
+    decay_mult: 1
+  }
+  param {
+    lr_mult: 2
+    decay_mult: 0
+  }
+  convolution_param {
+    num_output: 2
+    kernel_size: 1
+    stride: 1
+     weight_filler {
+      type: "xavier"
+    }
+    bias_filler {
+      type: "constant"
+      value: 0
+    }
+  }
+}
+
+layer {
+  name: "conv4-2"
+  type: "Convolution"
+  bottom: "conv3"
+  top: "conv4-2"
+  param {
+    lr_mult: 1
+    decay_mult: 1
+  }
+  param {
+    lr_mult: 2
+    decay_mult: 0
+  }
+  convolution_param {
+    num_output: 4
+    kernel_size: 1
+    stride: 1
+     weight_filler {
+      type: "xavier"
+	}
+    bias_filler {
+      type: "constant"
+      value: 0
+    }
+  }
+}
+layer {
+  name: "prob1"
+  type: "Softmax"
+  bottom: "conv4-1"
+  top: "prob1"
+}