solar_panels_detection.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rYWHh3Nv0oIv"
      },
      "source": [
        "# Solar Panels Detection"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nd2kUcrU05jO"
      },
      "source": [
        "Training and Detection of solar panels using YOLOv8 and the MAXAR WorldView-3 30 cm dataset (Germany region).\n",
        "\n",
        "The original dataset was split into train, validation, and test. This step can be done with the notebook *solar_dataset_preparation.ipynb*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D1Xf8Vyw1CQY"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9rdGyVTYs-rb",
        "outputId": "c792fbed-f80e-41d1-f0fd-dcb4d33c30fa"
      },
      "outputs": [],
      "source": [
        "!pip -q install ultralytics # Ignore if not using Colab"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MYGG8AIO1FAy"
      },
      "source": [
        "## Import Libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D672qZqJ0nR-"
      },
      "outputs": [],
      "source": [
        "from ultralytics import YOLO\n",
        "import os\n",
        "import torch\n",
        "import yaml\n",
        "import random\n",
        "import glob\n",
        "import cv2\n",
        "from matplotlib import pyplot as plt\n",
        "try:\n",
        "    from google.colab import files\n",
        "    from google.colab import drive\n",
        "except:\n",
        "    print(\"Some libraries were ignored.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgBAeBWu1IKP"
      },
      "source": [
        "## User Inputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xRUrW0Ob0nj1"
      },
      "outputs": [],
      "source": [
        "# Datasets folder\n",
        "# Folder where the train, validation and test datasets are saved\n",
        "datasets_folder = \"datasets\"\n",
        "\n",
        "# YOLO model hyperparameters\n",
        "# Target image size for training\n",
        "image_size = 416\n",
        "# Total number of training epochs\n",
        "epochs_num = 100\n",
        "# Batch size for training. -1 is auto.\n",
        "batch_size = -1\n",
        "# Minimum confidence threshold for detections\n",
        "conf_thr = 0.2\n",
        "# Intersection Over Union threshold\n",
        "iou_thr = 0.7"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QMWoMaDK1Myn"
      },
      "source": [
        "## Code"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JqytP9FG4Eme"
      },
      "source": [
        "### Train and Validation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "USiU0JAR4Y_v"
      },
      "source": [
        "Ignore the **Train and Validation** step if you already trained the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dr-74Jej0nvE"
      },
      "outputs": [],
      "source": [
        "# Create new folder (if does not exist) to store the datasets\n",
        "if not os.path.exists(datasets_folder):\n",
        "  os.mkdir(datasets_folder)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Check if GPU is available\n",
        "if torch.cuda.is_available():\n",
        "  print(\"GPU available.\")\n",
        "  torch.cuda.set_device(0)\n",
        "  device = 0 # Single GPU, it has support for multiple\n",
        "else:\n",
        "  print(\"GPU not available, using CPU.\")\n",
        "  print(\"On Colab, go to: Edit > Notebook Settings > Hardware Accelerator > Select T4 GPU\")\n",
        "  torch.cuda.set_device(-1)\n",
        "  device = \"cpu\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Mount drive\n",
        "try:\n",
        "  drive.mount(\"/content/drive\")\n",
        "  save_driver_folder = \"/content/drive/MyDrive/datasets\"\n",
        "except:\n",
        "  print(\"Unable to mount drive.\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7UMUfLZB6p5b",
        "outputId": "4fb4422d-dcc5-43d4-eeb6-52cdb0555be4"
      },
      "outputs": [],
      "source": [
        "# Copies the Train, Validation and Test saved on datasets folder from your Drive to this VM\n",
        "try:\n",
        "  # Train\n",
        "  !cp {save_driver_folder+\"/train.zip\"} {datasets_folder}\n",
        "\n",
        "  # Validation\n",
        "  !cp {save_driver_folder+\"/validation.zip\"} {datasets_folder}\n",
        "\n",
        "  # Test\n",
        "  !cp {save_driver_folder+\"/test.zip\"} {datasets_folder}\n",
        "except:\n",
        "  print(\"Copy of datasets from Drive ignored.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VnkhMJca0HKo"
      },
      "outputs": [],
      "source": [
        "# Unzip the datasets\n",
        "# Train\n",
        "!unzip -q -o {datasets_folder+\"/train.zip\"} -d {datasets_folder+\"/train\"}\n",
        "os.remove(datasets_folder+\"/train.zip\")\n",
        "\n",
        "# Validation\n",
        "!unzip -q -o {datasets_folder+\"/validation.zip\"} -d {datasets_folder+\"/validation\"}\n",
        "os.remove(datasets_folder+\"/validation.zip\")\n",
        "\n",
        "# Test\n",
        "!unzip -q -o {datasets_folder+\"/test.zip\"} -d {datasets_folder+\"/test\"}\n",
        "os.remove(datasets_folder+\"/test.zip\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0cVdau1L0fWF",
        "outputId": "396d8f04-01af-4586-d3c7-a79326647389"
      },
      "outputs": [],
      "source": [
        "# Load best pre-trained YOLOv8 model\n",
        "yv8x_model = YOLO(\"models/yolov8x.pt\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JDFowpSR7RTg"
      },
      "outputs": [],
      "source": [
        "# Generate yaml file\n",
        "yaml_path = os.path.join(datasets_folder, \"solar_panels_data.yaml\")\n",
        "yaml_data = {\"names\":[\"Solar panels\"], \"nc\":1, \"path\":os.path.abspath(\"datasets\"),\n",
        "              \"train\":\"train/images\", \"val\":\"validation/images\", \"test\":\"test/images\"}\n",
        "with open(yaml_path, 'w') as yaml_file:\n",
        "    yaml.dump(yaml_data, yaml_file, default_flow_style=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "CB1iJNAc47ld",
        "outputId": "d0264f2b-52d1-4c36-f7de-da39d2d95c32"
      },
      "outputs": [],
      "source": [
        "train_results = yv8x_model.train(data=yaml_path, epochs=epochs_num, batch=batch_size,\n",
        "                                 seed=42, imgsz=image_size, verbose=True, val=True, plots=True, device=device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XQ4qo77d48Wo"
      },
      "outputs": [],
      "source": [
        "# Copy best weights to Drive\n",
        "try:\n",
        "  !cp {\"runs/detect/train/weights/best.pt\"} {save_driver_folder}\n",
        "except:\n",
        "  print(\"Saving to drive ignored.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ls9fFbNhBYpJ"
      },
      "outputs": [],
      "source": [
        "# Download the best weights\n",
        "try:\n",
        "  files.download(\"runs/detect/train/weights/best.pt\")\n",
        "except:\n",
        "  print(\"Download to local machine ignored.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DJPdtzYb4Kgj"
      },
      "source": [
        "### Test"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iLT2AkU64KyF"
      },
      "outputs": [],
      "source": [
        "# Import best trained YOLO solar panels model\n",
        "try:\n",
        "  best_model = YOLO(\"runs/detect/train/weights/best.pt\")\n",
        "except:\n",
        "  print(\"Loaded from drive.\")\n",
        "  !cp {save_driver_folder+\"/solar_100epochs_-1batch_yolov8x_best.pt\"} {datasets_folder}\n",
        "  best_model = YOLO(datasets_folder+\"/solar_100epochs_-1batch_yolov8x_best.pt\")\n",
        "  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bEtcTeCj_wAB"
      },
      "outputs": [],
      "source": [
        "# Try the new model the Test dataset\n",
        "test_images_folder = os.path.join(datasets_folder, \"test\", \"images\")\n",
        "test_results = best_model.predict(source=test_images_folder, imgsz=image_size, conf=conf_thr,\n",
        "                                  iou=iou_thr, device=device, show_labels=False,\n",
        "                                  show_conf=False, save=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Open and show random predicted images\n",
        "predict_images = glob.glob(os.path.join(\"runs\", \"detect\", \"predict\", \"*.tif\"))\n",
        "rand_index = random.randint(0, len(predict_images)-1)\n",
        "predict_image_path = predict_images[rand_index]\n",
        "\n",
        "image_bgr = cv2.imread(predict_image_path)\n",
        "image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)\n",
        "plt.imshow(image_rgb)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.10.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}