{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "## Example Usage of UniFace Library for Face Alignment\n",
    "This guide demonstrates how to use the **UniFace** library for face detection and face alignment. Follow the steps below to set up and execute the example.\n",
    "\n",
    "## 1. Install UniFace\n",
    "Install the **UniFace** library using `pip`. The `-q` flag suppresses logs for a clean output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q uniface"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Import Required Libraries\n",
    "Import the necessary libraries for image processing, visualization and face alignment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from uniface import RetinaFace, face_alignment, draw_detections\n",
    "from uniface.constants import RetinaFaceWeights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- `cv2`: Used for image reading and processing.\n",
    "- `numpy`: Used for converting model outputs to numpy.\n",
    "- `matplotlib`: To display inference resulst\n",
    "- `RetinaFace`: The model class from the **UniFace** library.\n",
    "- `face_alignment`: A utility function for face alignment.\n",
    "- `draw_detections`: A utility function to draw bounding boxes and landmarks on the image.\n",
    "\n",
    "## 3. Initialize the RetinaFace Model\n",
    "Initialize the RetinaFace model with a lightweight pre-trained backbone and detection parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-03-26 11:44:08,755 - INFO - Initializing RetinaFace with model=RetinaFaceWeights.MNET_V2, conf_thresh=0.5, nms_thresh=0.4, pre_nms_topk=5000, post_nms_topk=750, dynamic_size=False, input_size=(640, 640)\n",
      "2025-03-26 11:44:08,755 - INFO - Model 'RetinaFaceWeights.MNET_V2' already exists at C:\\Users\\yakhyo\\.uniface\\models\\RetinaFaceWeights.MNET_V2.onnx\n",
      "2025-03-26 11:44:08,767 - INFO - Verified model weights located at: C:\\Users\\yakhyo/.uniface/models\\RetinaFaceWeights.MNET_V2.onnx\n",
      "2025-03-26 11:44:08,825 - INFO - Successfully initialized the model from C:\\Users\\yakhyo/.uniface/models\\RetinaFaceWeights.MNET_V2.onnx\n"
     ]
    }
   ],
   "source": [
    "# Initialize the RetinaFace model\n",
    "uniface_inference = RetinaFace(\n",
    "    model_name=RetinaFaceWeights.MNET_V2,   # Model name\n",
    "    conf_thresh=0.5,                        # Confidence threshold\n",
    "    pre_nms_topk=5000,                      # Pre-NMS Top-K detections\n",
    "    nms_thresh=0.4,                         # NMS IoU threshold\n",
    "    post_nms_topk=750                       # Post-NMS Top-K detections,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Load and perform inference\n",
    "Load set of input images to perform face detection and alignment, storing the results for visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Paths to the input images\n",
    "image_paths = [\n",
    "    \"../assets/test_images/image0.jpg\",\n",
    "    \"../assets/test_images/image1.jpg\",\n",
    "    \"../assets/test_images/image2.jpg\",\n",
    "    \"../assets/test_images/image3.jpg\",\n",
    "    \"../assets/test_images/image4.jpg\",\n",
    "]\n",
    "\n",
    "# Lists to store detection results and aligned images\n",
    "detection_images = []\n",
    "aligned_images = []\n",
    "original_images = []\n",
    "\n",
    "# Process each image\n",
    "for image_path in image_paths:\n",
    "    # Load the image\n",
    "    input_image = cv2.imread(image_path)\n",
    "    if input_image is None:\n",
    "        print(f\"Error: Could not read image from {image_path}\")\n",
    "        continue\n",
    "    \n",
    "    # Perform face detection\n",
    "    boxes, landmarks = uniface_inference.detect(input_image)\n",
    "\n",
    "    if len(landmarks) == 0:\n",
    "        print(f\"No face detected in {image_path}\")\n",
    "        continue\n",
    "    \n",
    "    # Draw detections on the image for visualization\n",
    "    bbox_image = input_image.copy()\n",
    "    draw_detections(bbox_image, (boxes, landmarks), vis_threshold=0.6)\n",
    "\n",
    "    # Align the first detected face\n",
    "    landmark_array = landmarks[0]\n",
    "    aligned_image, _ = face_alignment(input_image, landmark_array, image_size=112)\n",
    "    \n",
    "    # Convert images to RGB format for proper visualization\n",
    "    bbox_image = cv2.cvtColor(bbox_image, cv2.COLOR_BGR2RGB) \n",
    "    aligned_image = cv2.cvtColor(aligned_image, cv2.COLOR_BGR2RGB)\n",
    "    \n",
    "    # Store the processed images for visualization\n",
    "    original_images.append(input_image)\n",
    "    detection_images.append(bbox_image)\n",
    "    aligned_images.append(aligned_image)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Display inference results\n",
    "Visualization of face detection and alignment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1360224 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 15 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot images in a 2-row layout\n",
    "fig, axes = plt.subplots(3, len(image_paths), figsize=(15, 10))\n",
    "\n",
    "# Titles for each row\n",
    "row_titles = [\"Input Images\", \"Detection Results\", \"Face Alignment\"]\n",
    "\n",
    "# Populate the grid with images\n",
    "for row, images in enumerate([original_images, detection_images, aligned_images]):\n",
    "    for col, img in enumerate(images):\n",
    "        # Display each image in the grid\n",
    "        axes[row, col].imshow(img)\n",
    "        axes[row, col].axis(\"off\")  # Remove axes for cleaner visuals\n",
    "        \n",
    "        # Set row title on the first column of each row\n",
    "        if col == 0:\n",
    "            axes[row, col].set_title(row_titles[row], fontsize=12, loc=\"left\")\n",
    "            \n",
    "\n",
    "# Adjust layout to prevent overlap and display the plot\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
