BEGC2024_YOLO_based_Building_Segmentation.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ohOR85LG0hLU"
      },
      "source": [
        "# **IEEE Bigdata Cup 2024: Building extraction**\n",
        "\n",
        "**Author:** [Yi-Jie Wong](https://www.linkedin.com/in/wongyijie/)<br>\n",
        "**Challenge link:** [Kaggle](https://www.kaggle.com/competitions/building-extraction-generalization-2024/leaderboard)<br>\n",
        "**Date created:** 2024/07/10<br>\n",
        "**Last modified:** 2024/09/12<br>\n",
        "**Description:** Cross-City Generalizability of Instance Segmentation Model in a Nationwide Building Extraction Task"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kL9rYPxcUVrX"
      },
      "source": [
        "## **Step 1: Setup Dependencies and Dataset (will auto restart session once complete)**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qR9zmEZrJM9_"
      },
      "outputs": [],
      "source": [
        "# clone this repo\n",
        "!git clone https://github.com/yjwong1999/RSBuildingExtraction.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ca2uXo7Fpz3b"
      },
      "outputs": [],
      "source": [
        "# install ultralytics\n",
        "%pip install ultralytics==8.1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aNvzg2rhwXyR"
      },
      "outputs": [],
      "source": [
        "!pip3 install pycocotools requests click"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qipEv7Z-od8B"
      },
      "outputs": [],
      "source": [
        "!pip install opendatasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JbnyP870-v1K"
      },
      "outputs": [],
      "source": [
        "# https://github.com/opengeos/leafmap/blob/e35e0a75a125614244e5913755c50ec4f307bcab/docs/notebooks/74_map_tiles_to_geotiff.ipynb#L7\n",
        "# require reload after installation\n",
        "\n",
        "%pip install -U leafmap\n",
        "!pip install mercantile"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wYG93C_m_lq6"
      },
      "outputs": [],
      "source": [
        "!pip install geomet==1.1.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RW5m7h4PJM-B"
      },
      "outputs": [],
      "source": [
        "# !pip install geopandas"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xgX4RR9g-x45"
      },
      "outputs": [],
      "source": [
        "# prompt: restart colab session\n",
        "\n",
        "import os\n",
        "os.kill(os.getpid(), 9)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W8sUDEvb3cG-",
        "jp-MarkdownHeadingCollapsed": true
      },
      "source": [
        "## **Step 2: Download and Setup the Dataset**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eAAO-2txJM-B"
      },
      "outputs": [],
      "source": [
        "# Download the IEEE BEGC 2024 dataset\n",
        "\n",
        "%cd RSBuildingExtraction\n",
        "\n",
        "import opendatasets as od\n",
        "\n",
        "od.download(\"https://www.kaggle.com/competitions/building-extraction-generalization-2024/data\")\n",
        "\n",
        "%cd ../"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vDYjvK_tJM-B"
      },
      "outputs": [],
      "source": [
        "# Setup the IEEE BEGC2024 dataset into the necessary format\n",
        "\n",
        "%cd RSBuildingExtraction\n",
        "\n",
        "# run the code\n",
        "!python setup_data.py\n",
        "\n",
        "%cd ../"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Acoh-fFW68E0"
      },
      "source": [
        "## **Step 3: Extract from Microsoft Building Footprint (BF) Dataset (OPTIONAL)**\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u7Jh2qh37QGO"
      },
      "source": [
        "### Step 3.1 - Define our area of interest (AOI)\n",
        "\n",
        "We define our area of interest (or AOI) as a GeoJSON geometry, then use the `shapely` library to get the bounding box.</br>\n",
        "Go to [https://geojson.io](https://geojson.io), and find the prefered AOI. Draw a box around the AOI, and you will get the coordindates for the AOI region.<br>\n",
        "Please make sure the selected AOI is covered in the [Microsoft Building Footprint dataset](https://github.com/microsoft/GlobalMLBuildingFootprints/), s not all region is covered.<br>\n",
        "We provide the AOI we used for Redmond, Washington and Las Vegas, Nevada. </br>\n",
        "However, we recommend using the Las Vegas AOI, which is better for the training.\n",
        "\n",
        "**Note**: the coordinate reference system is EPSG:4326. The coordinate in expressed as (long, lat) format."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cSFM6JJM7C49"
      },
      "outputs": [],
      "source": [
        "import leafmap\n",
        "import pandas as pd\n",
        "import geopandas as gpd\n",
        "from shapely import geometry\n",
        "import mercantile\n",
        "from tqdm import tqdm\n",
        "import os, shutil\n",
        "import tempfile\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "\n",
        "from IPython.display import clear_output"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KtmC4FY67OxD"
      },
      "outputs": [],
      "source": [
        "# # The Selected AOI is around Redmond, Washington\n",
        "\n",
        "# # Geometry copied from https://geojson.io\n",
        "# aoi_geom = {\n",
        "#     \"coordinates\": [\n",
        "#         [\n",
        "#             [-122.16484503187519, 47.69090474454916],\n",
        "#             [-122.16484503187519, 47.6217555345674],\n",
        "#             [-122.06529607517405, 47.6217555345674],\n",
        "#             [-122.06529607517405, 47.69090474454916],\n",
        "#             [-122.16484503187519, 47.69090474454916],\n",
        "#         ]\n",
        "#     ],\n",
        "#     \"type\": \"Polygon\",\n",
        "# }\n",
        "# aoi_shape = geometry.shape(aoi_geom)\n",
        "# minx, miny, maxx, maxy = aoi_shape.bounds"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ROEVVIJ678uY"
      },
      "outputs": [],
      "source": [
        "# The Selected AOI is around Las Vegas, Nevada (recommended)\n",
        "\n",
        "# Geometry copied from https://geojson.io\n",
        "aoi_geom = {\n",
        "    \"coordinates\": [\n",
        "        [\n",
        "            [-115.31432742408262, 36.27297250862463],\n",
        "            [-115.31432742408262, 36.00372747612303],\n",
        "            [-114.98257204779121, 36.00372747612303],\n",
        "            [-114.98257204779121, 36.27297250862463],\n",
        "            [-115.31432742408262, 36.27297250862463],\n",
        "        ]\n",
        "    ],\n",
        "    \"type\": \"Polygon\",\n",
        "}\n",
        "aoi_shape = geometry.shape(aoi_geom)\n",
        "minx, miny, maxx, maxy = aoi_shape.bounds"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "07P3Sr4_7emO"
      },
      "source": [
        "### Step 3.2 - Determine which tiles intersect our AOI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "22g07lO87dtB"
      },
      "outputs": [],
      "source": [
        "quad_keys = set()\n",
        "for tile in list(mercantile.tiles(minx, miny, maxx, maxy, zooms=9)):\n",
        "    quad_keys.add(mercantile.quadkey(tile))\n",
        "quad_keys = list(quad_keys)\n",
        "print(f\"The input area spans {len(quad_keys)} tiles: {quad_keys}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "itnc-m3q7gEh"
      },
      "source": [
        "### Step 3.3 - Download the building footprints for each tile that intersects our AOI and crop the results\n",
        "\n",
        "This is where most of the magic happens. We download all the building footprints for each tile that intersects our AOI, then only keep the footprints that are _contained_ by our AOI.\n",
        "\n",
        "*Note*: this step might take awhile depending on how many tiles your AOI covers and how many buildings footprints are in those tiles."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ibUl3pQq7hXW"
      },
      "outputs": [],
      "source": [
        "df = pd.read_csv(\n",
        "    \"https://minedbuildings.blob.core.windows.net/global-buildings/dataset-links.csv\", dtype=str\n",
        ")\n",
        "df.head()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sRcSrPBy7i54"
      },
      "outputs": [],
      "source": [
        "# create an empty dataframe\n",
        "df_poly = pd.DataFrame()\n",
        "\n",
        "# Obtain polygons for each tile that intersects the input geometry\n",
        "for quad_key in tqdm(quad_keys):\n",
        "    rows = df[df[\"QuadKey\"] == quad_key]\n",
        "    if rows.shape[0] == 1:\n",
        "        url = rows.iloc[0][\"Url\"]\n",
        "\n",
        "        df2 = pd.read_json(url, lines=True)\n",
        "        df2[\"geometry\"] = df2[\"geometry\"].apply(geometry.shape)\n",
        "        df_poly = pd.concat([df_poly, df2], ignore_index=True)\n",
        "\n",
        "    elif rows.shape[0] > 1:\n",
        "        raise ValueError(f\"Multiple rows found for QuadKey: {quad_key}! We are not sure how to use such data, so feel free to contribute!\")\n",
        "    else:\n",
        "        raise ValueError(f\"QuadKey not found in dataset: {quad_key}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bLExNwD17kkL"
      },
      "source": [
        "### Step 3.4 - Get the outer bbox of AOI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2C9AcNn27lli"
      },
      "outputs": [],
      "source": [
        "df_poly_geometry = df_poly['geometry']\n",
        "gdf = gpd.GeoDataFrame(df_poly_geometry, crs=\"EPSG:4326\")\n",
        "_, _, maxx, maxy = gdf.bounds.max()\n",
        "minx, miny, _, _ = gdf.bounds.min()\n",
        "outer_aoi_bbox = [minx, miny, maxx, maxy]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2TqCfpZc7o3_"
      },
      "source": [
        "### Step 3.5 - Crop AOI into tiles and obtain outer bbox of each tile"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "clN7eKC77nsx"
      },
      "outputs": [],
      "source": [
        "def cropAOI(outer_aoi_bbox, step):\n",
        "    minx, miny, maxx, maxy = outer_aoi_bbox\n",
        "    maxx = (maxx - minx)//step *step + minx\n",
        "    maxy = (maxy - miny)//step *step + miny\n",
        "    outer_aoi_bbox = minx, miny, maxx, maxy\n",
        "\n",
        "    aoi_bbox_list = []\n",
        "    # handle large image situation, crop into tiles\n",
        "    if (maxx - minx) > step or (maxy - miny) > step:\n",
        "        new_minx, new_maxy = minx, maxy\n",
        "\n",
        "        num_x_tiles = int((maxx - minx)//step)\n",
        "        num_y_tiles = int((maxy - miny)//step)\n",
        "\n",
        "        print(f'Number of x tiles: {num_x_tiles}')\n",
        "        print(f'Number of y tiles: {num_y_tiles}')\n",
        "\n",
        "        print(f'\\nTotal number of tiles: {num_x_tiles*num_y_tiles}')\n",
        "        for i in range(num_y_tiles):\n",
        "            new_miny = new_maxy - step\n",
        "            for j in range(num_x_tiles):\n",
        "                new_maxx = new_minx + step\n",
        "\n",
        "                aoi_bbox = [new_minx, new_miny, new_maxx, new_maxy]\n",
        "                aoi_bbox_list.append(aoi_bbox)\n",
        "\n",
        "                new_minx = new_maxx\n",
        "\n",
        "            #     break\n",
        "\n",
        "            new_minx = minx\n",
        "            new_maxy = new_miny\n",
        "\n",
        "            # break\n",
        "\n",
        "    return aoi_bbox_list, outer_aoi_bbox"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aq_Mz5oQ7rUV"
      },
      "outputs": [],
      "source": [
        "step = 0.0009\n",
        "aoi_bbox_list, outer_aoi_bbox = cropAOI(outer_aoi_bbox, step)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q93gH0k47syD"
      },
      "source": [
        "### Step 3.6 - Label each tile sequentially"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yTppG4UX7tVT"
      },
      "outputs": [],
      "source": [
        "def labelAOITile(outer_aoi_bbox, bounded_df, step):\n",
        "    minx, miny, maxx, maxy = outer_aoi_bbox\n",
        "\n",
        "    bounded_df['left_tile'] = (bounded_df['minx'] - minx)//step\n",
        "    bounded_df['right_tile'] = (bounded_df['maxx'] - minx)//step + 1\n",
        "    bounded_df['bottom_tile'] = (bounded_df['miny'] - miny)//step - 1\n",
        "    bounded_df['top_tile'] = (bounded_df['maxy'] - miny)//step\n",
        "\n",
        "    num_x_tiles = (maxx - minx)//step\n",
        "    num_y_tiles = (maxy - miny)//step\n",
        "\n",
        "    bounded_df['Tile'] = bounded_df.apply(lambda row: calculateTile(row, num_x_tiles, num_y_tiles, step), axis=1)\n",
        "    # bounded_df = bounded_df.drop(['left_tile', 'right_tile', 'top_tile', 'bottom_tile'], axis=1)\n",
        "\n",
        "    return bounded_df\n",
        "\n",
        "def calculateTile(row, num_x_tiles, num_y_tiles, step):\n",
        "    y_diff = int(row['top_tile'] - row['bottom_tile'])\n",
        "    x_diff = int(row['right_tile'] - row['left_tile'])\n",
        "\n",
        "    y_buffer = num_y_tiles - row['top_tile']\n",
        "\n",
        "    tile_list = []\n",
        "    for i in range(y_diff):\n",
        "        for j in range(x_diff):\n",
        "            tile_list.append(int(row['left_tile'] + j + (i + y_buffer)*num_x_tiles))\n",
        "\n",
        "    return tile_list"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tTxUBaPt7vWI"
      },
      "outputs": [],
      "source": [
        "bounded_df = labelAOITile(outer_aoi_bbox, gdf.bounds, step)\n",
        "bounded_df.head()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2wNGHQ2G7vu0"
      },
      "outputs": [],
      "source": [
        "# outer_aoi_bbox"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_JyLCgaG7wnI"
      },
      "outputs": [],
      "source": [
        "# bounded_df[bounded_df['Tile'].apply(lambda x: 262144 in x)].head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GJdVTyJh7yUD"
      },
      "source": [
        "### Step 3.7 - Normalise Polygon and save image with labels"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fSJknSFs7xc5"
      },
      "outputs": [],
      "source": [
        "df_poly = pd.concat([bounded_df, df_poly_geometry], axis=1)\n",
        "df_poly.head()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZvfFDMtx70q5"
      },
      "outputs": [],
      "source": [
        "def normaliseBbox(aoi_bbox_list, df_poly, step, start_idx=0, end_idx=-1, shuffle=False):\n",
        "    tile_list = df_poly['Tile'].to_list()\n",
        "    tile_list_unpacked = []\n",
        "    for sublist in tile_list:\n",
        "        for item in sublist:\n",
        "            tile_list_unpacked.append(item)\n",
        "    tile_list = set(tile_list_unpacked)\n",
        "\n",
        "    print(f'Number of tiles to be generated: {len(tile_list)}')\n",
        "    print('\\nGenerating tiles...')\n",
        "\n",
        "    # convert\n",
        "    tile_list = list(tile_list)\n",
        "\n",
        "    # set np random seed and shuffle tile_list\n",
        "    if shuffle:\n",
        "        np.random.seed(0)\n",
        "        np.random.shuffle(tile_list)\n",
        "\n",
        "    tile_list = tile_list[start_idx:end_idx] # testing purpose\n",
        "\n",
        "    for tile in tqdm(tile_list):\n",
        "        try:\n",
        "            print(f'\\n~ Tile {tile}')\n",
        "            aoi_bbox = aoi_bbox_list[tile]\n",
        "            cropped_minx, cropped_miny, cropped_maxx, cropped_maxy = aoi_bbox\n",
        "\n",
        "            # df_tile = df_poly[df_poly['Tile'] == tile]\n",
        "            df_tile = df_poly[df_poly['Tile'].apply(lambda x: tile in x)]\n",
        "            for idx, row in df_tile.iterrows():\n",
        "                polygon = row['geometry']\n",
        "                # Extract coordinates and adjust them\n",
        "                new_exterior = []\n",
        "                for x, y in polygon.exterior.coords:\n",
        "                    # Adjust x and y coordinates based on the cropped bounding box\n",
        "                    new_x = (x - cropped_minx) / (cropped_maxx - cropped_minx)\n",
        "                    new_y = (y - cropped_miny) / (cropped_maxy - cropped_miny)\n",
        "\n",
        "                    if new_x < 0:\n",
        "                        new_x = 0\n",
        "                    elif new_x > 1:\n",
        "                        new_x = 1\n",
        "\n",
        "                    if new_y < 0:\n",
        "                        new_y = 0\n",
        "                    elif new_y > 1:\n",
        "                        new_y = 1\n",
        "\n",
        "                    # print(x - cropped_minx, x, cropped_minx, cropped_maxx, new_x)\n",
        "                    # print(y - cropped_miny, y, cropped_miny, cropped_maxy, new_y)\n",
        "\n",
        "                    new_exterior.append((new_x, new_y))\n",
        "\n",
        "                # Create a new polygon with adjusted coordinates\n",
        "                new_polygon = geometry.Polygon(new_exterior)\n",
        "\n",
        "                # Update the 'geometry' column in the DataFrame\n",
        "                df_tile.at[idx, 'geometry'] = new_polygon\n",
        "\n",
        "            df_tile['geometry_str'] = df_tile['geometry'].astype(str)\n",
        "            df_tile['geometry_str'] = df_tile['geometry_str'].str.lstrip('POLYGON ((').str.rstrip('))').str.replace(',', '')\n",
        "\n",
        "            # save segmentation labels\n",
        "            with open(f'{os.getcwd()}/labels/tile_{tile}.txt', 'w') as f:\n",
        "                for item in df_tile['geometry_str']:\n",
        "                    f.write(f'0 {item}\\n')\n",
        "\n",
        "            # save tif image\n",
        "            output_path = f'{os.getcwd()}/images_tiff/tile_{tile}.tif'\n",
        "            zoom = 20\n",
        "            leafmap.map_tiles_to_geotiff(output_path, aoi_bbox, zoom=zoom, source='SATELLITE')\n",
        "\n",
        "            clear_output()\n",
        "\n",
        "            # print(df_tile.head())\n",
        "\n",
        "        except Exception as e:\n",
        "            print(f'Error occurred for tile {tile}: {e}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "or958PIYBc4w"
      },
      "outputs": [],
      "source": [
        "image_dir_path = f'{os.getcwd()}/images_tiff'\n",
        "label_dir_path = f'{os.getcwd()}/labels'\n",
        "\n",
        "if os.path.exists(image_dir_path):\n",
        "    shutil.rmtree(image_dir_path)\n",
        "if os.path.exists(label_dir_path):\n",
        "    shutil.rmtree(label_dir_path)\n",
        "\n",
        "os.makedirs(image_dir_path)\n",
        "os.makedirs(label_dir_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "08fTnVpo72RJ"
      },
      "outputs": [],
      "source": [
        "start_idx = 0\n",
        "end_idx = 3000\n",
        "shuffle = True\n",
        "normaliseBbox(aoi_bbox_list, df_poly, step, start_idx, end_idx, shuffle)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6j5hvh_M73iZ"
      },
      "outputs": [],
      "source": [
        "# df_poly[df_poly['Tile'].apply(lambda x: 1 in x)].head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "44T5SJxD75Js",
        "jp-MarkdownHeadingCollapsed": true
      },
      "source": [
        "### Step 3.8 - Convert Tiff to JPEG"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0si4JYX974cD"
      },
      "outputs": [],
      "source": [
        "def convert_tiff_to_jpeg(input_dir, output_dir):\n",
        "    # check if output_dir exists, if not create it\n",
        "    if not os.path.exists(output_dir):\n",
        "        os.makedirs(output_dir)\n",
        "\n",
        "    for filename in os.listdir(input_dir):\n",
        "        # check if file is an image (ends with .tif)\n",
        "        if filename.endswith('.tif'):\n",
        "            img = Image.open(os.path.join(input_dir, filename))\n",
        "            img = img.transpose(Image.FLIP_TOP_BOTTOM)\n",
        "\n",
        "            # check if image is RGB mode, if not convert it\n",
        "            if img.mode != 'RGB':\n",
        "                img = img.convert('RGB')\n",
        "\n",
        "            # create new filename, replace .tif with .jpg\n",
        "            output_filename = os.path.splitext(filename)[0] + '.jpg'\n",
        "\n",
        "            # save the image in JPEG format\n",
        "            img.save(os.path.join(output_dir, output_filename), 'JPEG')\n",
        "\n",
        "    print(\"Conversion from TIFF to JPEG completed.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8BsiNT4KBzL5"
      },
      "outputs": [],
      "source": [
        "if os.path.exists(f'{os.getcwd()}/images_jpeg'):\n",
        "    shutil.rmtree(f'{os.getcwd()}/images_jpeg')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cmr7rwDD77Mt"
      },
      "outputs": [],
      "source": [
        "convert_tiff_to_jpeg(f'{os.getcwd()}/images_tiff/', f'{os.getcwd()}/images_jpeg/')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0YJtNpbV7-Hp",
        "jp-MarkdownHeadingCollapsed": true
      },
      "source": [
        "### Step 3.9 - Display Image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c21DM6sn79oO"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from matplotlib.patches import Polygon\n",
        "import matplotlib.pyplot as plt\n",
        "from PIL import Image\n",
        "\n",
        "i = 0\n",
        "image_dir = f'{os.getcwd()}/images_tiff/'\n",
        "img_name = sorted(os.listdir(image_dir))[i]\n",
        "label_name = img_name.replace('tif', 'txt')\n",
        "\n",
        "# Load the image\n",
        "im = Image.open(f'{os.getcwd()}/images_tiff/{img_name}')\n",
        "im = im.transpose(Image.FLIP_TOP_BOTTOM)\n",
        "im_array = np.array(im)\n",
        "\n",
        "# Create a figure and axes\n",
        "fig, ax = plt.subplots()\n",
        "\n",
        "# Display the image\n",
        "ax.imshow(im_array)\n",
        "\n",
        "width, height = im.size\n",
        "print(f'Image width: {width}')\n",
        "print(f'Image height: {height}')\n",
        "\n",
        "# Read and plot the polygons from geometries.txt\n",
        "with open(f'{os.getcwd()}/labels/{label_name}', 'r') as f:\n",
        "    for line in f.readlines():\n",
        "        coords_str = line.strip().split(' ')[1:]\n",
        "        coords = [float(c) for c in coords_str]\n",
        "        coords_x = coords[0::2]\n",
        "        coords_x = [x * width for x in coords_x]\n",
        "\n",
        "        coords_y = coords[1::2]\n",
        "        coords_y = [y * height for y in coords_y]\n",
        "        polygon = Polygon(list(zip(coords_x, coords_y)), closed=True, fill=False, edgecolor='r')\n",
        "        ax.add_patch(polygon)\n",
        "\n",
        "# Set the axis limits to match the image dimensions\n",
        "ax.set_xlim(0, im_array.shape[1])\n",
        "ax.set_ylim(im_array.shape[0], 0)  # Invert y-axis for image display\n",
        "\n",
        "# Show the plot\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OJXtGw8MNDex",
        "jp-MarkdownHeadingCollapsed": true
      },
      "source": [
        "### Step 3.10 - Copy to `RSBuildingExtraction/mydata` (the BEGC2024 training set in YOLO format)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "etcmZE7gNEef"
      },
      "outputs": [],
      "source": [
        "cwd = os.getcwd()\n",
        "!cp -r {cwd}/images_jpeg/* {cwd}/RSBuildingExtraction/mydata/images/train\n",
        "!cp -r {cwd}/labels/* {cwd}/RSBuildingExtraction/mydata/labels/train"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x = os.listdir('/content/RSBuildingExtraction/mydata/images/train')\n",
        "len(x)"
      ],
      "metadata": {
        "id": "tIK5oEXHNBeS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Step 3: Download our Preloaded Redmond/Las Vegas dataset (PREFERED)**"
      ],
      "metadata": {
        "id": "CSLt9qrfRhQZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import locale\n",
        "locale.getpreferredencoding = lambda: \"UTF-8\""
      ],
      "metadata": {
        "id": "ed165nCuhJtB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# # Download our Redmond dataset\n",
        "# !curl -L -o \"Redmond.zip\" \"https://www.dropbox.com/scl/fi/9d9yw2d3z777iiequ7mmr/Redmond.zip?rlkey=l9noi6ffibkmag76isteh6is1&st=lxl5ahd1&dl=0\"\n",
        "# !unzip \"Redmond.zip\"\n",
        "\n",
        "# # Copy to RSBuildingExtraction/mydata\n",
        "# cwd = os.getcwd()\n",
        "# !cp -r {cwd}/Redmond/images_jpeg/* {cwd}/RSBuildingExtraction/mydata/images/train\n",
        "# !cp -r {cwd}/Redmond/labels/* {cwd}/RSBuildingExtraction/mydata/labels/train"
      ],
      "metadata": {
        "id": "tQqCRZ85RToc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Download our Las Vegas dataset\n",
        "!curl -L -o \"LasVegas.zip\" \"https://www.dropbox.com/scl/fi/qaqo2lqd7x7gxh521f8ce/LasVegas.zip?rlkey=j3oute8e9ia9yoa1hc85fw2ev&st=h6agw363&dl=0\"\n",
        "!unzip \"LasVegas.zip\"\n",
        "\n",
        "# Copy to RSBuildingExtraction/mydata\n",
        "import os\n",
        "cwd = os.getcwd()\n",
        "!cp -r {cwd}/LasVegas/images_jpeg/* {cwd}/RSBuildingExtraction/mydata/images/train\n",
        "!cp -r {cwd}/LasVegas/labels/* {cwd}/RSBuildingExtraction/mydata/labels/train"
      ],
      "metadata": {
        "id": "ApS7HKqNQs5Y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "len(os.listdir('RSBuildingExtraction/mydata/images/train'))"
      ],
      "metadata": {
        "id": "CcXOWQ-iR00G"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "len(os.listdir('RSBuildingExtraction/mydata/labels/train'))"
      ],
      "metadata": {
        "id": "6gHibG_1SLLB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fqWw_CXugWyV"
      },
      "source": [
        "## **Step 4: Training YOLOv8-seg**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hkY7FJCA0Uh0"
      },
      "outputs": [],
      "source": [
        "from ultralytics import YOLO\n",
        "import os, shutil\n",
        "\n",
        "# yaml file of the dataset\n",
        "yaml_file = \"RSBuildingExtraction/mydata/data.yaml\"\n",
        "\n",
        "# use OBB pretrained YOLOv8 models for transfer learning\n",
        "model = YOLO(\"yolov8m-seg.pt\").load(\"yolov8m-obb.pt\")\n",
        "\n",
        "# Train the model (mainly shutdown mosaic + add flipud + add rotation)\n",
        "results = model.train(data=yaml_file, epochs=50, imgsz=640, plots=True, mixup=0.2)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# You can also download our pretrained weights here\n",
        "\n",
        "import locale\n",
        "locale.getpreferredencoding = lambda: \"UTF-8\"\n",
        "\n",
        "# You can also download our pretrained weights\n",
        "!curl -L -o \"yolov8m-seg_LasVegas.pt\" \"https://www.dropbox.com/scl/fi/cdrl62i3mx9p82lqwpik5/yolov8m-seg_LasVegas.pt?rlkey=8ao7a5zz7xnqfd74deffprix2&st=0k24i2xp&dl=0\""
      ],
      "metadata": {
        "id": "zuNiUC1CeW2x"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SAuZjm61KrDd"
      },
      "outputs": [],
      "source": [
        "# prompt: make prediction using model on all images in RSBuildingExtraction/building-extraction-generalization-2024/test/image using python\n",
        "\n",
        "import os\n",
        "from ultralytics import YOLO\n",
        "\n",
        "# Load the trained YOLOv8 model\n",
        "# model = YOLO('runs/segment/train/weights/last.pt') # newly trained from scratch\n",
        "model = YOLO('yolov8m-seg_LasVegas.pt') # our pretrained models\n",
        "\n",
        "# Directory containing test images\n",
        "test_image_dir = 'RSBuildingExtraction/building-extraction-generalization-2024/test/image'\n",
        "\n",
        "# Decoding according to the .yaml file class names order\n",
        "decoding_of_predictions ={0: 'building'}\n",
        "\n",
        "# Iterate through images in the test directory\n",
        "IDs = []\n",
        "entries = []\n",
        "for image_filename in sorted(os.listdir(test_image_dir)):\n",
        "    # remove extension from image_filename\n",
        "    ID = int(os.path.splitext(image_filename)[0])\n",
        "    print(ID)\n",
        "\n",
        "    image_path = os.path.join(test_image_dir, image_filename)\n",
        "\n",
        "    # Perform prediction on the image\n",
        "    results = model.predict(source=image_path, save=True, conf=0.2, imgsz=640, iou=0.9)\n",
        "\n",
        "    # Print results for each image (optional)\n",
        "    for r in results:\n",
        "        conf_list = r.boxes.conf.cpu().numpy().tolist()\n",
        "        clss_list = r.boxes.cls.cpu().numpy().tolist()\n",
        "        original_list = clss_list\n",
        "        updated_list = []\n",
        "        for element in original_list:\n",
        "                updated_list.append(decoding_of_predictions[int(element)])\n",
        "\n",
        "    # bounding_boxes = r.boxes.xyxy.cpu().numpy()\n",
        "\n",
        "    confidences = conf_list\n",
        "    class_names = updated_list\n",
        "    try:\n",
        "        masks = r.masks.xy\n",
        "    except:\n",
        "        masks = []\n",
        "\n",
        "    # Check if bounding boxes, confidences and class names match\n",
        "    if len(masks) != len(confidences) or len(masks) != len(class_names):\n",
        "        print(\"Error: Number of bounding boxes, confidences, and class names should be the same.\")\n",
        "        continue\n",
        "\n",
        "    entry = []\n",
        "    for m in masks:\n",
        "        temp = []\n",
        "        if len(m) <4:\n",
        "            continue\n",
        "        for xy in m:\n",
        "            x, y = xy[0], xy[1]\n",
        "            temp.append((int(x), int(y)))\n",
        "        entry.append(temp)\n",
        "\n",
        "    IDs.append(ID)\n",
        "    entries.append(entry)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3S7bhutH8lZL"
      },
      "outputs": [],
      "source": [
        "# prompt: create a csv with two columns, using list IDs and entries as the columns, add ImageID and Coordinates as the column title\n",
        "\n",
        "import csv\n",
        "\n",
        "# Assuming you have the 'IDs' and 'entries' lists as defined in the previous code\n",
        "\n",
        "# Create a list of dictionaries to store the data\n",
        "data = []\n",
        "for i in range(len(IDs)):\n",
        "  data.append({'ImageID': IDs[i], 'Coordinates': entries[i]})\n",
        "\n",
        "# Write the data to a CSV file\n",
        "with open('output.csv', 'w', newline='') as csvfile:\n",
        "  fieldnames = ['ImageID', 'Coordinates']\n",
        "  writer = csv.DictWriter(csvfile, fieldnames=fieldnames)\n",
        "\n",
        "  writer.writeheader()\n",
        "  writer.writerows(data)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "92hzWdGSZKUM"
      },
      "outputs": [],
      "source": [
        "# zip output\n",
        "\n",
        "!zip output.zip output.csv\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iDcRm3EMPdYV"
      },
      "outputs": [],
      "source": [
        "!kaggle competitions submit -c building-extraction-generalization-2024 -f \"/content/output.csv\" -m \"YOLOv8m-seg with Las Vegas dataset\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gBcdO25GECos"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bvx4g8L_HwRS"
      },
      "source": [
        "## **Step 5: Visualization**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hEIvAf_GEuXB"
      },
      "outputs": [],
      "source": [
        "import cv2\n",
        "import matplotlib.pyplot as plt\n",
        "from ultralytics import YOLO\n",
        "import numpy as np\n",
        "\n",
        "# Load the trained YOLOv8 model\n",
        "# model = YOLO('runs/segment/train/weights/last.pt') # newly trained from scratch\n",
        "model = YOLO('yolov8m-seg_LasVegas.pt') # our pretrained models\n",
        "\n",
        "# index\n",
        "idx = 998 # 90 # 50 # 510 # 0 # 1 # 257 # 998 # 600 # 700 # 701 # 705 # 751 # 754 # 125 # 15\n",
        "idx_formatted = str(idx).zfill(4)\n",
        "\n",
        "# path\n",
        "test_img_path = f'RSBuildingExtraction/building-extraction-generalization-2024/test/image/{idx_formatted}.tif'\n",
        "test_output_path = f'runs/segment/predict/{idx_formatted}.tif'\n",
        "\n",
        "# Create a figure and axes\n",
        "fig, axs = plt.subplots(1, 4, figsize=(20, 5))\n",
        "\n",
        "# The rest of your plotting code would go here, using axs[0], axs[1], etc. to access the individual subplots.\n",
        "img = plt.imread(test_img_path)\n",
        "axs[0].imshow(img)\n",
        "axs[0].set_title('Original Image')\n",
        "\n",
        "# for loop different iou\n",
        "for i, iou in enumerate([0.7, 0.8, 0.9]):\n",
        "    # Perform prediction on the image\n",
        "    results = model.predict(source=test_img_path, save=False, conf=0.2, imgsz=640, iou=iou)\n",
        "\n",
        "    # Create binary mask\n",
        "    b_mask = np.zeros(img.shape[:2], np.uint8)\n",
        "\n",
        "    # mask\n",
        "    try:\n",
        "        masks = results[0].masks.xy\n",
        "    except:\n",
        "        masks = []\n",
        "\n",
        "    # loop all mask\n",
        "    for m in masks:\n",
        "        # if less than 4 points, ignore\n",
        "        if len(m) <4:\n",
        "            continue\n",
        "        # get contour\n",
        "        contour = m.astype(np.int32)\n",
        "        contour = contour.reshape(-1, 1, 2)\n",
        "        _ = cv2.drawContours(b_mask, [contour], -1, (255, 255, 255), cv2.FILLED)\n",
        "\n",
        "    axs[i + 1].imshow(b_mask, cmap='gray')\n",
        "    axs[i + 1].set_title(f'NMS IoU Threshold = {iou}')\n",
        "\n",
        "# Remove axis ticks\n",
        "for ax in axs:\n",
        "    ax.set_xticks([])\n",
        "    ax.set_yticks([])\n",
        "\n",
        "# Show the plot\n",
        "# plt.tight_layout()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6XhxwrwFEUPm"
      },
      "outputs": [],
      "source": [
        "import cv2\n",
        "import matplotlib.pyplot as plt\n",
        "from ultralytics import YOLO\n",
        "import numpy as np\n",
        "\n",
        "# index\n",
        "idx = 998 # 90 # 50 # 510 # 0 # 1 # 257 # 998 # 600 # 700 # 701 # 705 # 751 # 754 # 125 # 15\n",
        "idx_formatted = str(idx).zfill(4)\n",
        "\n",
        "# path\n",
        "test_img_path = f'RSBuildingExtraction/building-extraction-generalization-2024/test/image/{idx_formatted}.tif'\n",
        "test_output_path = f'runs/segment/predict/{idx_formatted}.tif'\n",
        "\n",
        "# Create a figure and axes\n",
        "fig, axs = plt.subplots(1, 4, figsize=(20, 5))\n",
        "\n",
        "# The rest of your plotting code would go here, using axs[0], axs[1], etc. to access the individual subplots.\n",
        "img = plt.imread(test_img_path)\n",
        "axs[0].imshow(img)\n",
        "axs[0].set_title('Original Image')\n",
        "\n",
        "# for loop different iou\n",
        "weights = ['yolov8m-seg_basic.pt', 'yolov8m-seg_Washington.pt', 'yolov8m-seg_LasVegas.pt']\n",
        "names = [\"YOLOv8m Basic\", \"YOLOv8m (w' Washington Dataset)\", \"YOLOv8m (w' Las Vegas Dataset)\"]\n",
        "for i, weight in enumerate(weights):\n",
        "    # Load the trained YOLOv8 model\n",
        "    model = YOLO(weight)\n",
        "\n",
        "    # Perform prediction on the image\n",
        "    results = model.predict(source=test_img_path, save=False, conf=0.2, imgsz=640, iou=0.9)\n",
        "\n",
        "    # Create binary mask\n",
        "    b_mask = np.zeros(img.shape[:2], np.uint8)\n",
        "\n",
        "    # mask\n",
        "    try:\n",
        "        masks = results[0].masks.xy\n",
        "    except:\n",
        "        masks = []\n",
        "\n",
        "    # loop all mask\n",
        "    for m in masks:\n",
        "        # if less than 4 points, ignore\n",
        "        if len(m) <4:\n",
        "            continue\n",
        "        # get contour\n",
        "        contour = m.astype(np.int32)\n",
        "        contour = contour.reshape(-1, 1, 2)\n",
        "        _ = cv2.drawContours(b_mask, [contour], -1, (255, 255, 255), cv2.FILLED)\n",
        "\n",
        "    axs[i + 1].imshow(b_mask, cmap='gray')\n",
        "    axs[i + 1].set_title(f'{names[i]}')\n",
        "\n",
        "# Remove axis ticks\n",
        "for ax in axs:\n",
        "    ax.set_xticks([])\n",
        "    ax.set_yticks([])\n",
        "\n",
        "# Show the plot\n",
        "# plt.tight_layout()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6n5dJ7RHT3eG"
      },
      "outputs": [],
      "source": [
        "import cv2\n",
        "import matplotlib.pyplot as plt\n",
        "from ultralytics import YOLO\n",
        "import numpy as np\n",
        "\n",
        "# index\n",
        "idx = 998 # 90 # 50 # 510 # 0 # 1 # 257 # 998 # 600 # 700 # 701 # 705 # 751 # 754 # 125 # 15\n",
        "idx_formatted = str(idx).zfill(4)\n",
        "\n",
        "# path\n",
        "test_img_path = f'RSBuildingExtraction/building-extraction-generalization-2024/test/image/{idx_formatted}.tif'\n",
        "test_output_path = f'runs/segment/predict/{idx_formatted}.tif'\n",
        "\n",
        "# Create a figure and axes\n",
        "fig, axs = plt.subplots(1, 4, figsize=(20, 5))\n",
        "\n",
        "# The rest of your plotting code would go here, using axs[0], axs[1], etc. to access the individual subplots.\n",
        "img = plt.imread(test_img_path)\n",
        "axs[0].imshow(img)\n",
        "axs[0].set_title('Original Image')\n",
        "\n",
        "# for loop different iou\n",
        "weights = ['yolov8n-seg_basic.pt', 'yolov8s-seg_basic.pt', 'yolov8m-seg_basic.pt']\n",
        "names = [\"YOLOv8n-seg\", \"YOLOv8s-seg\", \"YOLOv8m-seg\"]\n",
        "for i, weight in enumerate(weights):\n",
        "    # Load the trained YOLOv8 model\n",
        "    model = YOLO(weight)\n",
        "\n",
        "    # Perform prediction on the image\n",
        "    results = model.predict(source=test_img_path, save=False, conf=0.2, imgsz=640, iou=0.9)\n",
        "\n",
        "    # Create binary mask\n",
        "    b_mask = np.zeros(img.shape[:2], np.uint8)\n",
        "\n",
        "    # mask\n",
        "    try:\n",
        "        masks = results[0].masks.xy\n",
        "    except:\n",
        "        masks = []\n",
        "\n",
        "    # loop all mask\n",
        "    for m in masks:\n",
        "        # if less than 4 points, ignore\n",
        "        if len(m) <4:\n",
        "            continue\n",
        "        # get contour\n",
        "        contour = m.astype(np.int32)\n",
        "        contour = contour.reshape(-1, 1, 2)\n",
        "        _ = cv2.drawContours(b_mask, [contour], -1, (255, 255, 255), cv2.FILLED)\n",
        "\n",
        "    axs[i + 1].imshow(b_mask, cmap='gray')\n",
        "    axs[i + 1].set_title(f'{names[i]}')\n",
        "\n",
        "# Remove axis ticks\n",
        "for ax in axs:\n",
        "    ax.set_xticks([])\n",
        "    ax.set_yticks([])\n",
        "\n",
        "# Show the plot\n",
        "# plt.tight_layout()\n",
        "plt.show()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "kL9rYPxcUVrX",
        "W8sUDEvb3cG-",
        "Acoh-fFW68E0",
        "CSLt9qrfRhQZ",
        "fqWw_CXugWyV",
        "Bvx4g8L_HwRS"
      ],
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.10.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}