Train_Yolov8_OBB_Object_Detection_On_Custom_Dataset.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/woodsj1206/Train-Yolov8-OBB-Object-Detection-On-Custom-Dataset/blob/main/Train_Yolov8_OBB_Object_Detection_On_Custom_Dataset.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Step 1: Checking GPU Availability\n",
        "Check if a GPU is available in your Google Colab environment. A GPU can significantly accelerate the training process of deep learning models like YOLOv8."
      ],
      "metadata": {
        "id": "b0y2vHRgUyFG"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5S75OM8X7IH7"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Step 2: Unzipping Dataset Files\n",
        "Unzip the dataset files containing images and annotations. These files are typically compressed to save space and are now being extracted for use in training the YOLOv8 model.\n",
        "\n",
        "**NOTE:** Ensure that the images and their corresponding annotation files have the same names (e.g., image_1.png should have an associated image_1.txt file)."
      ],
      "metadata": {
        "id": "-EcCLxEyWAot"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Unzip the zip file containing the images\n",
        "!unzip -q 'REPLACE_WITH_YOUR_IMAGE_ZIP_FILE_PATH' -d '/content/images'\n",
        "\n",
        "# Unzip the zip file containing the annotations\n",
        "!unzip -q 'REPLACE_WITH_YOUR_ANNOTATIONS_ZIP_FILE_PATH' -d '/content/annotations'"
      ],
      "metadata": {
        "id": "szt5QH1vSwR_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Step 3: Installing Ultralytics\n",
        "Install Ultralytics, a library that simplifies working with YOLO object detection models."
      ],
      "metadata": {
        "id": "qoScN3sjXQmt"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_5AMobU32Z-E"
      },
      "outputs": [],
      "source": [
        "# Install Ultralytics library\n",
        "!pip install ultralytics\n",
        "\n",
        "# Import necessary modules from Ultralytics\n",
        "from ultralytics import YOLO"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Step 4: Mounting Google Drive\n",
        "Mount Google Drive to the Colab environment, enabling access to files stored in your Google Drive."
      ],
      "metadata": {
        "id": "eFjeGMT4X27V"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dob4YRhec9Fm"
      },
      "outputs": [],
      "source": [
        "# Import the drive module from google.colab\n",
        "from google.colab import drive\n",
        "\n",
        "# Mount Google Drive to '/content/Google_Drive'\n",
        "drive.mount('/content/Google_Drive')"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Step 5: Define Root Directory and Create Subdirectories\n",
        "Define the root directory for your project and create necessary subdirectories to organize your data."
      ],
      "metadata": {
        "id": "01fjuyyHZ98E"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xBHfolzSdnR5"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# Prompt user for the path\n",
        "PATH = input(\"Enter the desired path: \")\n",
        "\n",
        "ROOT_DIR = f'/content/Google_Drive/MyDrive/{PATH}'\n",
        "\n",
        "# Create directory for data\n",
        "DATA_DIR = os.path.join(ROOT_DIR, 'data')\n",
        "\n",
        "# Create directory for images\n",
        "IMAGES_DIR = os.path.join(DATA_DIR, 'images')\n",
        "IMAGES_VAL_DIR = os.path.join(IMAGES_DIR, 'val')\n",
        "IMAGES_TRAIN_DIR = os.path.join(IMAGES_DIR, 'train')\n",
        "\n",
        "# Create directory for labels\n",
        "LABELS_DIR = os.path.join(DATA_DIR, 'labels')\n",
        "LABELS_VAL_DIR = os.path.join(LABELS_DIR, 'val')\n",
        "LABELS_TRAIN_DIR = os.path.join(LABELS_DIR, 'train')\n",
        "\n",
        "# Create directory for testing\n",
        "TESTING_DIR = os.path.join(ROOT_DIR, 'testing')\n",
        "\n",
        "# Create the directories if the root directory doesn't exist\n",
        "if not os.path.exists(ROOT_DIR):\n",
        "    os.makedirs(ROOT_DIR)\n",
        "\n",
        "    os.makedirs(DATA_DIR)\n",
        "\n",
        "    os.makedirs(IMAGES_DIR)\n",
        "    os.makedirs(IMAGES_VAL_DIR)\n",
        "    os.makedirs(IMAGES_TRAIN_DIR)\n",
        "\n",
        "    os.makedirs(LABELS_DIR)\n",
        "    os.makedirs(LABELS_VAL_DIR)\n",
        "    os.makedirs(LABELS_TRAIN_DIR)\n",
        "\n",
        "    os.makedirs(TESTING_DIR)\n",
        "\n",
        "    print(f\"Root directory '{ROOT_DIR}' created successfully.\")\n",
        "else:\n",
        "    print(f\"Root directory '{ROOT_DIR}' already exists.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Step 6: Organize Dataset for Training and Validation\n",
        "Organize the dataset by moving images and corresponding annotations into separate directories for training and validation."
      ],
      "metadata": {
        "id": "MjnBrJwNpaZv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import shutil\n",
        "\n",
        "# Replace placeholders with actual paths to your images and annotations\n",
        "IMAGES_PATH = 'REPLACE_WITH_PATH_TO_YOUR_IMAGES'\n",
        "ANNOTATIONS_PATH = 'REPLACE_WITH_PATH_TO_YOUR_ANNOTATIONS'\n",
        "\n",
        "# Get list of image and annotation files\n",
        "image_files = os.listdir(IMAGES_PATH)\n",
        "annotation_files = os.listdir(ANNOTATIONS_PATH)\n",
        "\n",
        "# Sort the files for consistency\n",
        "image_files.sort()\n",
        "annotation_files.sort()\n",
        "\n",
        "# Determine the number of files for the training set\n",
        "train_count = int(len(image_files) * 0.7) # 70% of data used for training, 30% used for validation\n",
        "\n",
        "# Move files to the training directory\n",
        "for file in image_files[:train_count]:\n",
        "    shutil.move(os.path.join(IMAGES_PATH, file), os.path.join(IMAGES_TRAIN_DIR, file))\n",
        "for file in annotation_files[:train_count]:\n",
        "    shutil.move(os.path.join(ANNOTATIONS_PATH, file), os.path.join(LABELS_TRAIN_DIR, file))\n",
        "\n",
        "# Move files to the validation directory\n",
        "for file in image_files[train_count:]:\n",
        "    shutil.move(os.path.join(IMAGES_PATH, file), os.path.join(IMAGES_VAL_DIR, file))\n",
        "for file in annotation_files[train_count:]:\n",
        "    shutil.move(os.path.join(ANNOTATIONS_PATH, file), os.path.join(LABELS_VAL_DIR, file))"
      ],
      "metadata": {
        "id": "yo3wB7Ueai8s"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Step 7: Generate YAML Configuration File\n",
        "Create a YAML configuration file specifying the paths to the training and validation datasets, as well as the class names (labels) used in your dataset."
      ],
      "metadata": {
        "id": "DzrPODTrqLb2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import yaml\n",
        "\n",
        "# Define the data structure for the YAML file\n",
        "data = {\n",
        "    'path': f'{DATA_DIR}',\n",
        "    'train': 'images/train',\n",
        "    'val': 'images/val',\n",
        "    'names': {\n",
        "        # Add your labels here\n",
        "        0: 'REPLACE_WITH_YOUR_LABEL_NAME'\n",
        "      # 1: ''\n",
        "      # 2: ''\n",
        "      # ...\n",
        "    }\n",
        "}\n",
        "\n",
        "# Construct the full path to the YAML file\n",
        "output_file = os.path.join(ROOT_DIR, \"config.yaml\")\n",
        "\n",
        "# Write the data to the YAML file\n",
        "with open(output_file, 'w') as yaml_file:\n",
        "    yaml.dump(data, yaml_file)"
      ],
      "metadata": {
        "id": "dJAsMnHG7YL7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Step 8: Train the YOLOv8 Model\n",
        "Train the YOLOv8 model using the provided dataset configuration and save the training results."
      ],
      "metadata": {
        "id": "Dn93_5WkqiNm"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oyZJX6PVfE7J"
      },
      "outputs": [],
      "source": [
        "# Load a pre-trained YOLOv8 model\n",
        "model = YOLO(\"yolov8n-obb.pt\")\n",
        "\n",
        "# Train the model using the provided dataset configuration\n",
        "model_results = model.train(data=os.path.join(ROOT_DIR, \"config.yaml\"), epochs=20)\n",
        "\n",
        "# Save the training results\n",
        "shutil.make_archive(base_dir='/content/runs', root_dir='/content/runs', format='zip', base_name=f'{ROOT_DIR}/runs')"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Step 9: View Training Results\n",
        "Display the training results and the confusion matrix generated during the training process."
      ],
      "metadata": {
        "id": "oJO_S4RYq4at"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gVZypTYk6T85"
      },
      "outputs": [],
      "source": [
        "from IPython.display import Image\n",
        "\n",
        "# Display the training results\n",
        "Image('/content/runs/obb/train/results.png')"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Display the confusion matrix\n",
        "Image('/content/runs/obb/train/confusion_matrix.png')"
      ],
      "metadata": {
        "id": "Hv6luQVEAVHK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Step 10: Process Video Files\n",
        "Process video files using the trained YOLOv8 model and save the processed videos with bounding boxes and class labels.\n",
        "\n",
        "**NOTE:** Ensure you have uploaded your video files to the testing folder before proceeding with this step."
      ],
      "metadata": {
        "id": "SUFhvf6LsUIt"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import cv2\n",
        "\n",
        "# Define the directory where you want to save the images\n",
        "OUTPUT_DIR = '/content/output'\n",
        "\n",
        "# Create the output directory if it doesn't exist\n",
        "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
        "\n",
        "# Path to the trained model file\n",
        "model_file = \"/content/runs/obb/train/weights/best.pt\"\n",
        "\n",
        "# Load the trained model\n",
        "model = YOLO(model_file)\n",
        "\n",
        "# List all files in the folder\n",
        "video_files = os.listdir(TESTING_DIR)\n",
        "\n",
        "# Process each video file\n",
        "for video_file in video_files:\n",
        "    # Open the video file\n",
        "    cap = cv2.VideoCapture(os.path.join(TESTING_DIR, video_file))\n",
        "\n",
        "    # Get the dimensions and frame rate of the video\n",
        "    w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))\n",
        "\n",
        "    # Create the output video writer\n",
        "    out = cv2.VideoWriter(f'{OUTPUT_DIR}/{video_file}_OBB_object_detection.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))\n",
        "\n",
        "    # Read the first frame of the video\n",
        "    ret, frame = cap.read()\n",
        "\n",
        "    while ret:\n",
        "        # Perform object detection on the frame\n",
        "        results = model(frame)[0]\n",
        "\n",
        "        for result in results.obb.data.tolist():\n",
        "            x_center, y_center, width, height, angle, score, class_id = result\n",
        "\n",
        "            # Calculate the coordinates of the bounding box corners\n",
        "            corners = cv2.boxPoints(((x_center, y_center), (width, height), angle * 180 / np.pi))\n",
        "            corners = np.int0(corners)\n",
        "\n",
        "            # Draw the bounding box\n",
        "            cv2.drawContours(frame, [corners], 0, (255, 0, 0), 2)\n",
        "\n",
        "            # Display class name and score\n",
        "            label = f'{results.names[int(class_id)]}: {score:.2f}'\n",
        "\n",
        "            # Get the size of the text\n",
        "            text_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)[0]\n",
        "\n",
        "            # Draw a filled rectangle behind the text label\n",
        "            cv2.rectangle(frame, (int(x_center) - 5, int(y_center) - text_size[1] - 10), (int(x_center) + text_size[0] + 5, int(y_center) + 5), (255, 0, 0), cv2.FILLED)\n",
        "\n",
        "            # Draw the text label\n",
        "            cv2.putText(frame, label, (int(x_center), int(y_center) - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)\n",
        "\n",
        "        out.write(frame)\n",
        "        ret, frame = cap.read()\n",
        "\n",
        "    # Release the video capture and writer objects\n",
        "    cap.release()\n",
        "    out.release()\n",
        "    cv2.destroyAllWindows()"
      ],
      "metadata": {
        "id": "0YojVvRieRw3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Step 11: Display Processed Video\n",
        "Use the MoviePy library to display the processed video within the Colab environment."
      ],
      "metadata": {
        "id": "2QRTh2G_smo0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from moviepy.editor import *\n",
        "\n",
        "# Path to the processed video file\n",
        "path = 'REPLACE_WITH_PATH_TO_VIDEO_FILE'\n",
        "\n",
        "# Load the processed video clip\n",
        "clip = VideoFileClip(path)\n",
        "\n",
        "# Display the video clip\n",
        "clip.ipython_display(height=540, width=960)"
      ],
      "metadata": {
        "id": "j3W2ukMNRN6B"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}