FaceEnhancement.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DqUjdIu2KjUx"
      },
      "source": [
        "# Install required libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7NLKy062DxEw"
      },
      "outputs": [],
      "source": [
        "!pip install opencv-python torch torchvision"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nhWe43OoKzm1"
      },
      "outputs": [],
      "source": [
        "import cv2\n",
        "import os\n",
        "\n",
        "def extract_frames(video_source, output_folder, save_every_n_frames=5):\n",
        "    \"\"\"\n",
        "    Extracts frames from a video and saves every nth frame to the specified folder.\n",
        "\n",
        "    Args:\n",
        "        video_source (str): Path to the video file.\n",
        "        output_folder (str): Directory to save the extracted frames.\n",
        "        save_every_n_frames (int): Save every nth frame (default is 5).\n",
        "    \"\"\"\n",
        "    os.makedirs(output_folder, exist_ok=True)\n",
        "\n",
        "    cap = cv2.VideoCapture(video_source)\n",
        "\n",
        "    if not cap.isOpened():\n",
        "        print(f\"Error opening video source: {video_source}\")\n",
        "        return\n",
        "\n",
        "    # Process video frames\n",
        "    frame_number = 0\n",
        "    while True:\n",
        "        ret, frame = cap.read()\n",
        "        if not ret:\n",
        "            break\n",
        "\n",
        "        # Save every nth frame\n",
        "        if frame_number % save_every_n_frames == 0:\n",
        "            frame_filename = os.path.join(output_folder, f'frame_{frame_number}.jpg')\n",
        "            cv2.imwrite(frame_filename, frame)\n",
        "            print(f\"Saved: {frame_filename}\")\n",
        "\n",
        "        frame_number += 1\n",
        "\n",
        "    cap.release()\n",
        "\n",
        "    print(\"Frame extraction complete.\")\n",
        "\n",
        "# Example usage\n",
        "video_source = '/content/1643-148614430_small.mp4'\n",
        "output_folder = '/content/1643_extracted_frames-new'\n",
        "extract_frames(video_source, output_folder, save_every_n_frames=5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LLMIytzpRJjs"
      },
      "source": [
        "# Use MTCNN with YOLOv5(underlying model) for face detection in frames"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-BhdDicARIVV"
      },
      "outputs": [],
      "source": [
        "!pip install facenet_pytorch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q0HK7hfQRSiF"
      },
      "outputs": [],
      "source": [
        "import cv2\n",
        "import torch\n",
        "from facenet_pytorch import MTCNN\n",
        "from PIL import Image\n",
        "import os\n",
        "\n",
        "# Initialize the MTCNN model\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "mtcnn = MTCNN(keep_all=True, device=device)\n",
        "\n",
        "# Directory to save the extracted frames\n",
        "output_folder = '/content/1643_extracted_frames_highlighted_faces'\n",
        "os.makedirs(output_folder, exist_ok=True)\n",
        "\n",
        "\n",
        "def process_frame(frame, frame_number):\n",
        "    try:\n",
        "        # Convert OpenCV image (BGR) to PIL image (RGB)\n",
        "        pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))\n",
        "\n",
        "        # Detect faces\n",
        "        boxes, probs = mtcnn.detect(pil_image)\n",
        "\n",
        "        if boxes is not None:\n",
        "            for box in boxes:\n",
        "\n",
        "                x1, y1, x2, y2 = map(int, box)\n",
        "                width = x2 - x1\n",
        "                height = y2 - y1\n",
        "\n",
        "\n",
        "                top_margin = int(height * 0.5)\n",
        "                side_margin = int(width * 0.3)\n",
        "\n",
        "                # Modify bounding box\n",
        "                x1 = max(0, x1 - side_margin)\n",
        "                y1 = max(0, y1 - top_margin)\n",
        "                x2 = min(frame.shape[1], x2 + side_margin)\n",
        "                y2 = min(frame.shape[0], y2 + side_margin)\n",
        "\n",
        "                # Draw bounding boxes on the frame\n",
        "                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)\n",
        "\n",
        "        return frame\n",
        "    except Exception as e:\n",
        "        print(f\"Error processing frame: {e}\")\n",
        "        return frame\n",
        "\n",
        "# Open video file\n",
        "video_source = '/content/1643-148614430_small.mp4'  # Replace with your video file path\n",
        "cap = cv2.VideoCapture(video_source)\n",
        "\n",
        "# Check if video source is opened successfully\n",
        "if not cap.isOpened():\n",
        "    print(f\"Error opening video source: {video_source}\")\n",
        "    exit()\n",
        "\n",
        "# Process video frames\n",
        "frame_number = 0\n",
        "save_every_n_frames = 5  # Save every 5th frame\n",
        "while True:\n",
        "    ret, frame = cap.read()\n",
        "    if not ret:\n",
        "        break\n",
        "\n",
        "    # Process the frame and highlight faces\n",
        "    frame = process_frame(frame, frame_number)\n",
        "\n",
        "    # Save every 5th frame with highlighted faces\n",
        "    if frame_number % save_every_n_frames == 0:\n",
        "        frame_filename = os.path.join(output_folder, f'frame_{frame_number}.jpg')\n",
        "        cv2.imwrite(frame_filename, frame)\n",
        "        print(f\"Saved: {frame_filename}\")\n",
        "\n",
        "    frame_number += 1\n",
        "\n",
        "# Release resources\n",
        "cap.release()\n",
        "\n",
        "print(\"Frame extraction complete.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "erG4TJcsR1UX"
      },
      "outputs": [],
      "source": [
        "# frame extraction no detection\n",
        "import cv2\n",
        "import os\n",
        "\n",
        "# Directory to save the extracted frames\n",
        "output_folder = '/content/1643_extracted_frames-new'\n",
        "os.makedirs(output_folder, exist_ok=True)\n",
        "\n",
        "# Open video file\n",
        "video_source = '/content/1643-148614430_small.mp4'  # Replace with your video file path\n",
        "cap = cv2.VideoCapture(video_source)\n",
        "\n",
        "# Check if video source is opened successfully\n",
        "if not cap.isOpened():\n",
        "    print(f\"Error opening video source: {video_source}\")\n",
        "    exit()\n",
        "\n",
        "# Process video frames\n",
        "frame_number = 0\n",
        "save_every_n_frames = 5  # Save every 5th frame\n",
        "while True:\n",
        "    ret, frame = cap.read()\n",
        "    if not ret:\n",
        "        break\n",
        "\n",
        "    # Save every 5th frame\n",
        "    if frame_number % save_every_n_frames == 0:\n",
        "        frame_filename = os.path.join(output_folder, f'frame_{frame_number}.jpg')\n",
        "        cv2.imwrite(frame_filename, frame)\n",
        "        print(f\"Saved: {frame_filename}\")\n",
        "\n",
        "    frame_number += 1\n",
        "\n",
        "# Release resources\n",
        "cap.release()\n",
        "\n",
        "print(\"Frame extraction complete.\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bGMT7BYnVtJD"
      },
      "source": [
        "# Using GFPGAN for face enhancement"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nWEuz2TCV3g3"
      },
      "outputs": [],
      "source": [
        "!git clone https://github.com/TencentARC/GFPGAN.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qSkstjEZV4UW"
      },
      "outputs": [],
      "source": [
        "!pip install basicsr facexlib"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qCjsCnqZV4Rg"
      },
      "outputs": [],
      "source": [
        "!pip install -r /content/GFPGAN/requirements.txt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zLOpu23SV4O6"
      },
      "outputs": [],
      "source": [
        "# Change to the GFPGAN directory\n",
        "%cd /content/GFPGAN\n",
        "\n",
        "# List files to confirm VERSION file is present\n",
        "!ls\n",
        "\n",
        "# Run the setup.py script\n",
        "!python setup.py develop\n",
        "!pip install --use-pep517\n",
        "!pip install -r requirements.txt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OLh6ifunV4LU"
      },
      "outputs": [],
      "source": [
        "!wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth -P experiments/pretrained_models"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KPO-3-QSV4DI"
      },
      "outputs": [],
      "source": [
        "!python /content/GFPGAN/inference_gfpgan.py -i /content/1643_extracted_frames-new -o /content/enhanced_faces -v 1.3 -s 2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "03eC-0psWu6i"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E79hVtDGPilB"
      },
      "source": [
        "# Install the CodeFormer model for face Enhancement\n",
        "Link: [CodeFormer](https://shangchenzhou.com/projects/CodeFormer/)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9yvYFnwpPiJc"
      },
      "outputs": [],
      "source": [
        "!git clone https://github.com/sczhou/CodeFormer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vewYUSBEP5Kg"
      },
      "outputs": [],
      "source": [
        "!pip install basicsr facexlib"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "piguGZmbS-jp"
      },
      "source": [
        "### run the next cell only if error related to torch occurs\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BIRnPQPZVQp2"
      },
      "outputs": [],
      "source": [
        "!pip uninstall torch torchvision -y\n",
        "!pip install torch==1.13.0 torchvision==0.14.0 -f https://download.pytorch.org/whl/cu121/torch_stable.html"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BOWLmcgqPsWw"
      },
      "outputs": [],
      "source": [
        "!pip install -r /content/CodeFormer/requirements.txt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BnskMwFkPuQ6"
      },
      "outputs": [],
      "source": [
        "%cd /content/CodeFormer\n",
        "\n",
        "!python basicsr/setup.py develop"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vsi1Wi8mTmeG"
      },
      "outputs": [],
      "source": [
        "!pip install ninja ffmpeg dlib"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4IJydJMwTm48"
      },
      "outputs": [],
      "source": [
        "%cd /content/CodeFormer\n",
        "!python scripts/download_pretrained_models.py CodeFormer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-XiUnkL2Twm1"
      },
      "outputs": [],
      "source": [
        "%cd /content/CodeFormer\n",
        "\n",
        "!python inference_codeformer.py -w 0.5 --input_path /content/1643_extracted_frames-new --output_path /content/enhanced_faces"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oX1n04__VxSI"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}