Deforestation_inference.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c529015a",
   "metadata": {},
   "source": [
    "### **For given aoi, run Deforestation model(Unet-Diff) inference**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fed86eeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "import json\n",
    "import logging\n",
    "import os\n",
    "import re\n",
    "import shutil\n",
    "import tempfile\n",
    "import urllib\n",
    "import uuid\n",
    "import warnings\n",
    "from datetime import datetime, timedelta\n",
    "import time\n",
    "from os.path import join\n",
    "from pathlib import Path\n",
    "\n",
    "import cv2\n",
    "import geojson\n",
    "import geopandas\n",
    "import geopandas as gp\n",
    "import geopandas as gpd\n",
    "import imageio\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pyproj\n",
    "import rasterio\n",
    "import rasterio.mask\n",
    "import segmentation_models_pytorch as smp\n",
    "import shapely\n",
    "import torch\n",
    "import torch.backends.cudnn as cudnn\n",
    "from geojson import Feature\n",
    "from geopandas import GeoSeries\n",
    "from google.oauth2 import service_account\n",
    "from googleapiclient.discovery import build\n",
    "from googleapiclient.http import MediaIoBaseDownload\n",
    "from rasterio import Affine, features\n",
    "from rasterio.mask import mask as riomask\n",
    "from rasterio.merge import merge\n",
    "from rasterio.plot import reshape_as_image, reshape_as_raster\n",
    "from rasterio.warp import Resampling, calculate_default_transform, reproject\n",
    "from rasterio.windows import Window\n",
    "from scipy import spatial\n",
    "from sentinel2download.downloader import Sentinel2Downloader, logger\n",
    "from sentinel2download.overlap import Sentinel2Overlap\n",
    "from shapely import wkt\n",
    "from shapely.geometry import MultiPolygon, Polygon, box\n",
    "from shapely.ops import transform, unary_union\n",
    "from skimage.exposure import match_histograms\n",
    "from torch import nn\n",
    "from torchvision import transforms\n",
    "from tqdm import tqdm\n",
    "from rasterio.mask import mask as rasterio_mask\n",
    "import shapely\n",
    "\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "\n",
    "# Functions to create, load the model(Use them instead of Catalyst lib code)\n",
    "def prepare_device():\n",
    "    return torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "\n",
    "def prepare_model(model):\n",
    "    device = prepare_device()\n",
    "\n",
    "    if torch.cuda.is_available():\n",
    "        cudnn.benchmark = True\n",
    "\n",
    "    if torch.cuda.device_count() > 1:\n",
    "        raise ValueError(\"Multi GPU mode is not supported\")\n",
    "    else:\n",
    "        model = model.to(device)\n",
    "\n",
    "    return model, device\n",
    "\n",
    "\n",
    "def load_model(network, model_weights_path):\n",
    "    if network == \"unet-diff\":\n",
    "        aux_params = dict(\n",
    "            pooling=\"max\",  # one of 'avg', 'max'\n",
    "            dropout=0.1,  # dropout ratio, default is None\n",
    "            activation=\"sigmoid\",  # activation function, default is None\n",
    "            classes=1,  # define number of output labels\n",
    "        )\n",
    "        model = smp.Unet(\n",
    "            \"resnet18\",\n",
    "            aux_params=aux_params,\n",
    "            encoder_weights=None,\n",
    "            in_channels=27,\n",
    "            encoder_depth=2,\n",
    "            decoder_channels=(256, 128),\n",
    "        )\n",
    "\n",
    "    else:\n",
    "        raise ValueError(\"Unknown network\")\n",
    "\n",
    "    model, device = prepare_model(model)\n",
    "    model.load_state_dict(\n",
    "        torch.load(model_weights_path, map_location=torch.device(device))\n",
    "    )\n",
    "    return model, device"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cceca116",
   "metadata": {},
   "source": [
    "### 0. Setting up parameters\n",
    "#### Read Input from environment and setup folders and filenames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60529d0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "current_time = time.strftime(\"deforestation_cache_%Y_%m_%d_%H:%M:%S\", time.gmtime())\n",
    "REQUEST_ID = os.getenv(\"REQUEST_ID\")\n",
    "START_DATE = os.getenv(\"START_DATE\")\n",
    "\n",
    "END_DATE = os.getenv(\"END_DATE\")\n",
    "AOI = os.getenv(\"AOI\")\n",
    "SENTINEL2_GOOGLE_API_KEY = os.getenv(\"SENTINEL2_GOOGLE_API_KEY\")\n",
    "SATELLITE_CACHE_FOLDER = os.path.join(os.getenv(\"SENTINEL2_CACHE\"), current_time)\n",
    "\n",
    "INPUT_FOLDER = os.path.dirname(SATELLITE_CACHE_FOLDER)\n",
    "OUTPUT_FOLDER = os.getenv(\"OUTPUT_FOLDER\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9b4fe66",
   "metadata": {},
   "outputs": [],
   "source": [
    "PREPARED_DATA_FOLDER = os.path.join(SATELLITE_CACHE_FOLDER, \"prepared\")\n",
    "CODE_FOLDER = \"/code\"\n",
    "\n",
    "LANDCOVER_POLYGONS_PATH = os.path.join(CODE_FOLDER, \"data\", \"landcovers\")\n",
    "LANDCOVER_FILENAME = (\n",
    "    \"S2A_OPER_GIP_TILPAR_MPC__20151209T095117_V20150622T000000_21000101T000000_B00.kml\"\n",
    ")\n",
    "SENTINEL_TILES = os.path.join(LANDCOVER_POLYGONS_PATH, LANDCOVER_FILENAME)\n",
    "\n",
    "SCOPES = [\"https://www.googleapis.com/auth/drive.file\"]\n",
    "\n",
    "CLOUD_DATA_FOLDER = os.path.join(CODE_FOLDER, \"data\", \"clouds\")\n",
    "MODEL_PATH = os.path.join(CODE_FOLDER, \"models\", \"unet_diff.pth\")\n",
    "\n",
    "LOAD_DIR = SATELLITE_CACHE_FOLDER\n",
    "\n",
    "PRODUCT_TYPE = \"L1C\"\n",
    "BANDS = {\"TCI\", \"B01\", \"B02\", \"B04\", \"B05\", \"B08\", \"B8A\", \"B09\", \"B10\", \"B11\", \"B12\"}\n",
    "CONSTRAINTS = {\n",
    "    \"NODATA_PIXEL_PERCENTAGE\": 15.0,\n",
    "    \"CLOUDY_PIXEL_PERCENTAGE\": 15.0,\n",
    "}\n",
    "\n",
    "NEAREST_POLYGONS_NUMBER = 10\n",
    "# DATES_FOR_TILE = 2\n",
    "CLOUDS_PROBABILITY_THRESHOLD = 0.95\n",
    "REMOVE_OTHER_DATES = True\n",
    "\n",
    "NODATA = 0\n",
    "MAX_SHIFT_ITERS = 2\n",
    "MAX_SHIFT = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbf2ed2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(INPUT_FOLDER, exist_ok=True)\n",
    "os.makedirs(LANDCOVER_POLYGONS_PATH, exist_ok=True)\n",
    "os.makedirs(PREPARED_DATA_FOLDER, exist_ok=True)\n",
    "os.makedirs(CLOUD_DATA_FOLDER, exist_ok=True)\n",
    "os.makedirs(SATELLITE_CACHE_FOLDER, exist_ok=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1e125cce",
   "metadata": {},
   "source": [
    "### 1. Transform AOI to GeoJSON file and then download data from Sentinel L1C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d996771f",
   "metadata": {},
   "outputs": [],
   "source": [
    "aoi = gp.GeoDataFrame(geometry=[wkt.loads(AOI)], crs=\"epsg:4326\")\n",
    "aoi_filename = \"provided_aoi.geojson\"\n",
    "aoi.to_file(aoi_filename, driver=\"GeoJSON\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "494ed325",
   "metadata": {},
   "outputs": [],
   "source": [
    "s2overlap = Sentinel2Overlap(aoi_path=aoi_filename)\n",
    "overlap_tiles = s2overlap.overlap_with_geometry()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "533905b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def shift_date(date, delta=5, format=\"%Y-%m-%d\"):\n",
    "    date = datetime.strptime(date, format)\n",
    "    date = date - timedelta(days=delta)\n",
    "    return datetime.strftime(date, format)\n",
    "\n",
    "\n",
    "def diff_date(date_a, date_b, format=\"%Y-%m-%d\"):\n",
    "    date_a, date_b = datetime.strptime(date_a, format), datetime.strptime(\n",
    "        date_b, format\n",
    "    )\n",
    "    delta = date_a - date_b\n",
    "    return delta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21993ea9",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_timestamp, end_timestamp = datetime.strptime(\n",
    "    START_DATE, \"%Y-%m-%d\"\n",
    "), datetime.strptime(END_DATE, \"%Y-%m-%d\")\n",
    "year = ((end_timestamp - start_timestamp) / 2 + start_timestamp).year\n",
    "year"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ad90b82",
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_folder(folder, limit=15, nodata=0):\n",
    "    for file in os.listdir(folder):\n",
    "        if file.endswith(\".jp2\"):\n",
    "            with rasterio.open(os.path.join(folder, file)) as src:\n",
    "                data = src.read(1)\n",
    "                file_nodata_percentage = (\n",
    "                    round((data == nodata).sum() / data.size, 2) * 100\n",
    "                )\n",
    "            if file_nodata_percentage > limit:\n",
    "                return False\n",
    "            return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c77c7a5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_images(tiles, start_date, end_date):\n",
    "    loader = Sentinel2Downloader(SENTINEL2_GOOGLE_API_KEY)\n",
    "    loadings = dict()\n",
    "\n",
    "    for tile in tiles:\n",
    "        start = start_date\n",
    "        end = end_date\n",
    "\n",
    "        print(f\"Loading images for tile: {tile}...\")\n",
    "        count = 0\n",
    "        while count < MAX_SHIFT_ITERS:\n",
    "            loaded = loader.download(\n",
    "                PRODUCT_TYPE,\n",
    "                [tile],\n",
    "                start_date=start,\n",
    "                end_date=end,\n",
    "                output_dir=LOAD_DIR,\n",
    "                bands=BANDS,\n",
    "                constraints=CONSTRAINTS,\n",
    "                cores=2\n",
    "            )\n",
    "            \n",
    "            distinct_folders = set()\n",
    "            for duo_files in loaded:\n",
    "                distinct_folders.add(os.path.dirname(duo_files[0]))\n",
    "\n",
    "            loaded_ = []\n",
    "            for folder in distinct_folders:\n",
    "                if len(os.listdir(folder)) >= len(BANDS) + 1: #skip folder if it didn't have all needed BANDS\n",
    "                    if check_folder(\n",
    "                        folder, limit=CONSTRAINTS[\"NODATA_PIXEL_PERCENTAGE\"], nodata=NODATA\n",
    "                    ): # found if folders has nodata\n",
    "                        for duo_files in loaded:\n",
    "                            if folder in duo_files[0]:\n",
    "                                loaded_.append(duo_files)\n",
    "            loaded = loaded_\n",
    "            if not loaded:\n",
    "                previous_end = end\n",
    "                previous_start = start\n",
    "                end = start\n",
    "                start = shift_date(start, delta=MAX_SHIFT)\n",
    "                print(\n",
    "                    f\"For tile: {tile} and dates {previous_start} {previous_end} proper images not found! Shift dates to {start} {end}!\"\n",
    "                )\n",
    "            else:\n",
    "                break\n",
    "            count += 1\n",
    "        if loaded:\n",
    "            loadings[tile] = loaded\n",
    "            print(f\"Loading images for tile {tile} finished\")\n",
    "        else:\n",
    "            print(f\"Images for tile {tile} were not loaded!\")\n",
    "\n",
    "    # Remove MSK_QUALIT_* and MSK_DETFOO_* files for L1C tiles\n",
    "    removed_msk_loadings = {}\n",
    "    for tile, tile_paths in loadings.items():\n",
    "        removed_msk_loadings[tile] = []\n",
    "        for tile_path in tile_paths:\n",
    "            if \"MSK_\" not in tile_path[0]:\n",
    "                removed_msk_loadings[tile].append(tile_path)\n",
    "            else:\n",
    "                os.remove(tile_path[0])\n",
    "\n",
    "    return removed_msk_loadings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c95bcfd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "loadings_start_date = load_images(\n",
    "    overlap_tiles.Name.values, shift_date(START_DATE, delta=10), START_DATE\n",
    ")\n",
    "loadings_end_date = load_images(\n",
    "    overlap_tiles.Name.values, shift_date(END_DATE, delta=10), END_DATE\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ebe3059",
   "metadata": {},
   "outputs": [],
   "source": [
    "loadings_start_date"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22b35a43",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not loadings_start_date or not loadings_end_date:  \n",
    "    raise ValueError(\"Images not loaded for given AOI. Change dates, constraints\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ffcd624",
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_by_date(loadings, func=max, filtered=dict(), tag=\"start\"):\n",
    "    def _find_agg_date(folders, func=func):\n",
    "        dates = list()\n",
    "        for i, folder in enumerate(folders):\n",
    "            search = re.search(r\"_(\\d+)T\\d+_\", str(folder))\n",
    "            date = search.group(1)\n",
    "            date = datetime.strptime(date, \"%Y%m%d\")\n",
    "            dates.append(date)\n",
    "        last_date = func(dates)\n",
    "        last_date = datetime.strftime(last_date, \"%Y%m%d\")\n",
    "        return last_date\n",
    "\n",
    "    def _get_folder(files, last_date):\n",
    "        for idx, duo in enumerate(files):\n",
    "            if last_date in duo[0]:\n",
    "                return os.path.dirname(duo[0])\n",
    "        raise Exception(\"last_date folder not found\")\n",
    "\n",
    "    for tile, items in loadings.items():\n",
    "        try:\n",
    "            last_date = _find_agg_date(items)\n",
    "            bands_paths = dict()\n",
    "            for path, _ in items:\n",
    "                if PRODUCT_TYPE == \"L2A\":\n",
    "                    if last_date in path:\n",
    "                        if \"B8A_20m.jp2\" in path:\n",
    "                            bands_paths[\"B8A\"] = path\n",
    "                        if \"B11_20m.jp2\" in path:\n",
    "                            bands_paths[\"B11\"] = path\n",
    "                        if \"B04_10m.jp2\" in path:\n",
    "                            bands_paths[\"B04\"] = path\n",
    "                        if \"B08_10m.jp2\" in path:\n",
    "                            bands_paths[\"B08\"] = path\n",
    "                        if \"B12_20m.jp2\" in path:\n",
    "                            bands_paths[\"B12\"] = path\n",
    "                        if \"TCI_10m.jp2\" in path:\n",
    "                            bands_paths[\"TCI\"] = path\n",
    "                        folder = _get_folder(items, last_date)\n",
    "                elif PRODUCT_TYPE == \"L1C\":\n",
    "                    if last_date in path:\n",
    "                        if \"B8A.jp2\" in path:\n",
    "                            bands_paths[\"B8A\"] = path\n",
    "                        if \"B11.jp2\" in path:\n",
    "                            bands_paths[\"B11\"] = path\n",
    "                        if \"B04.jp2\" in path:\n",
    "                            bands_paths[\"B04\"] = path\n",
    "                        if \"B08.jp2\" in path:\n",
    "                            bands_paths[\"B08\"] = path\n",
    "                        if \"B12.jp2\" in path:\n",
    "                            bands_paths[\"B12\"] = path\n",
    "                        if \"TCI.jp2\" in path:\n",
    "                            bands_paths[\"TCI\"] = path\n",
    "                        folder_path = _get_folder(items, last_date)\n",
    "\n",
    "            info_dict = {\n",
    "                tag: dict(paths=bands_paths, date=last_date, folder=folder_path)\n",
    "            }\n",
    "\n",
    "            if tile in filtered.keys():\n",
    "                filtered[tile].update(info_dict)\n",
    "            else:\n",
    "                filtered.update({tile: info_dict})\n",
    "\n",
    "        except Exception as ex:\n",
    "            print(f\"Error for {tile}: {str(ex)}\")\n",
    "    return filtered"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f94524a",
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered = filter_by_date(loadings_start_date, func=max, tag=\"start\")\n",
    "filtered = filter_by_date(loadings_end_date, func=max, filtered=filtered, tag=\"end\")\n",
    "filtered"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e5e49cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_tiles_folders_dates(fitered):\n",
    "    tiles = tuple(fitered.keys())\n",
    "    start_date_folders = [filtered[tile][\"start\"][\"folder\"] for tile in tiles]\n",
    "    end_date_folders = [filtered[tile][\"end\"][\"folder\"] for tile in tiles]\n",
    "    dates = [filtered[tiles[0]]['start']['date'], filtered[tiles[0]]['end']['date']]\n",
    "    return tiles, start_date_folders, end_date_folders, dates\n",
    "\n",
    "\n",
    "tiles, start_date_folders, end_date_folders, dates = get_tiles_folders_dates(filtered)\n",
    "assert len(start_date_folders) == len(end_date_folders), 'Not enough tiles folders'\n",
    "start_date_folders, end_date_folders, tiles, dates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e7806bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "\n",
    "\n",
    "def remove_not_used_dates(\n",
    "    start_date_folder, end_date_folder, cache_dir=SATELLITE_CACHE_FOLDER\n",
    "):\n",
    "    start_split, end_split = start_date_folder.split(\"/\"), end_date_folder.split(\"/\")\n",
    "    if cache_dir != os.path.join(\"/\", *start_split[:-1]) or cache_dir != os.path.join(\n",
    "        \"/\", *end_split[:-1]\n",
    "    ):\n",
    "        raise ValueError(\"cache_dir is not valid\")\n",
    "    used_dates = start_split[-1], end_split[-1]\n",
    "    for folder in os.listdir(cache_dir):\n",
    "        if folder not in used_dates:\n",
    "            shutil.rmtree(os.path.join(cache_dir, folder))\n",
    "\n",
    "\n",
    "# if REMOVE_OTHER_DATES:\n",
    "#     remove_not_used_dates(start_date_folder, end_date_folder)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4dcc3036",
   "metadata": {},
   "source": [
    "### 2. Preparing images (calculating ndmi ndvi, scaling, merging to tiff)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8609d9a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def crop_aoi_from_raster(raster_path, aoi_gdf, output_path):\n",
    "    with rasterio.open(raster_path, 'r') as src:\n",
    "        meta = src.meta\n",
    "        aoi_gdf = aoi_gdf.to_crs(meta['crs'])\n",
    "        shapes_to_crop = [shapely.ops.unary_union(aoi_gdf.geometry.to_list())]\n",
    "        mask, out_trans = rasterio_mask(src,\n",
    "                                        shapes_to_crop,\n",
    "                                        nodata=0,\n",
    "                                        crop=True)\n",
    "\n",
    "        meta['height'] = mask.shape[1]\n",
    "        meta['width'] = mask.shape[2]\n",
    "        meta['transform'] = out_trans\n",
    "\n",
    "\n",
    "    with rasterio.open(output_path, 'w', **meta) as dst:\n",
    "        dst.write(mask)\n",
    "    print(f'saved raster {output_path}')\n",
    "    os.remove(raster_path)\n",
    "    print(f'removed raster {raster_path}')\n",
    "    return output_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bfc8bb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rasterio.merge import merge\n",
    "\n",
    "\n",
    "def transform_crs(data_path, save_path, dst_crs=\"EPSG:4326\", resolution=(10, 10)):\n",
    "    with rasterio.open(data_path) as src:\n",
    "        if resolution is None:\n",
    "            transform, width, height = calculate_default_transform(\n",
    "                src.crs, dst_crs, src.width, src.height, *src.bounds\n",
    "            )\n",
    "        else:\n",
    "            transform, width, height = calculate_default_transform(\n",
    "                src.crs,\n",
    "                dst_crs,\n",
    "                src.width,\n",
    "                src.height,\n",
    "                *src.bounds,\n",
    "                resolution=resolution,\n",
    "            )\n",
    "        kwargs = src.meta.copy()\n",
    "        kwargs.update(\n",
    "            {\"crs\": dst_crs, \"transform\": transform, \"width\": width, \"height\": height}\n",
    "        )\n",
    "        with rasterio.open(save_path, \"w\", **kwargs) as dst:\n",
    "            for i in range(1, src.count + 1):\n",
    "                reproject(\n",
    "                    source=rasterio.band(src, i),\n",
    "                    destination=rasterio.band(dst, i),\n",
    "                    src_transform=src.transform,\n",
    "                    src_crs=src.crs,\n",
    "                    dst_transform=transform,\n",
    "                    dst_crs=dst_crs,\n",
    "                    resampling=Resampling.nearest,\n",
    "                )\n",
    "\n",
    "    return save_path\n",
    "\n",
    "\n",
    "def stitch_tiles(paths, out_raster_path='test.tif'):\n",
    "    tiles = []\n",
    "    tmp_files = []\n",
    "    \n",
    "    for i, path in enumerate(paths):\n",
    "        if i == 0:\n",
    "            file = rasterio.open(path)\n",
    "            meta, crs = file.meta, file.crs\n",
    "        else:\n",
    "            tmp_path = path.replace(\n",
    "                '.jp2', '_tmp.jp2').replace('.tif', '_tmp.tif')\n",
    "            crs_transformed = transform_crs(path, tmp_path, \n",
    "                                            dst_crs=crs, \n",
    "                                            resolution=None)\n",
    "            tmp_files.append(crs_transformed)\n",
    "            file = rasterio.open(crs_transformed)\n",
    "        tiles.append(file)\n",
    "            \n",
    "    tile_arr, transform = rasterio.merge.merge(tiles, method='last')\n",
    "    \n",
    "    \n",
    "    meta.update({\"driver\": \"GTiff\",\n",
    "                 \"height\": tile_arr.shape[1],\n",
    "                 \"width\": tile_arr.shape[2],\n",
    "                 \"transform\": transform,\n",
    "                 \"crs\": crs})\n",
    "    \n",
    "    if '.jp2' in out_raster_path:\n",
    "        out_raster_path = out_raster_path.replace('.jp2', '_merged.tif')\n",
    "    else:\n",
    "        out_raster_path = out_raster_path.replace('.tif', '_merged.tif')\n",
    "    print(f'saved raster {out_raster_path}')\n",
    "\n",
    "    for tile in tiles:\n",
    "        tile.close()\n",
    "        \n",
    "    for tmp_file in tmp_files:\n",
    "        try:\n",
    "            os.remove(tmp_file)\n",
    "        except FileNotFoundError:\n",
    "            print(f'Tile {tmp_file} was removed or renamed, skipping')\n",
    "        \n",
    "    with rasterio.open(out_raster_path, \"w\", **meta) as dst:\n",
    "        dst.write(tile_arr)\n",
    "    \n",
    "    return out_raster_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27292af0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from time_dependent.data_prepare.prepare_tif import (get_ndmi, get_ndvi, merge,\n",
    "                                                     scale_img, search_band,\n",
    "                                                     to_tiff)\n",
    "\n",
    "\n",
    "def prepare_data(data_folder, save_path):\n",
    "    img_folder = data_folder\n",
    "\n",
    "    tmp_file = os.path.basename(data_folder)\n",
    "\n",
    "    os.makedirs(save_path, exist_ok=True)\n",
    "    save_file_merged = join(save_path, f\"all_merged_{tmp_file}.tif\")\n",
    "    print(save_file_merged)\n",
    "    bands, band_names = [\"TCI\", \"B08\", \"B8A\", \"B11\", \"B12\"], []\n",
    "\n",
    "    for band in bands:\n",
    "        band_names.append(join(img_folder, search_band(band, img_folder, \"jp2\")))\n",
    "\n",
    "    b4_name = join(img_folder, search_band(\"B04\", img_folder, \"jp2\"))\n",
    "    ndvi_name = join(img_folder, \"ndvi\")\n",
    "    ndmi_name = join(img_folder, \"ndmi\")\n",
    "    print(\"\\nall bands are converting to *tif...\\n\")\n",
    "\n",
    "    for band_name in band_names:\n",
    "        print(band_name[-3:])\n",
    "        if \"B08\" in band_name:\n",
    "            b8_name = band_name\n",
    "        if \"B8A\" in band_name:\n",
    "            b8a_name = band_name\n",
    "        if \"B11\" in band_name:\n",
    "            b11_name = band_name\n",
    "        to_tiff(f\"{band_name}.jp2\")\n",
    "\n",
    "    to_tiff(f\"{b4_name}.jp2\")\n",
    "    print(\"\\nndvi band is processing...\")\n",
    "\n",
    "    get_ndvi(f\"{b4_name}.tif\", f\"{b8_name}.tif\", f\"{ndvi_name}.tif\")\n",
    "\n",
    "    print(\"\\nndmi band is processing...\")\n",
    "\n",
    "    get_ndmi(f\"{b11_name}.tif\", f\"{b8a_name}.tif\", f\"{ndmi_name}.tif\")\n",
    "\n",
    "    band_names.append(ndvi_name)\n",
    "    band_names.append(ndmi_name)\n",
    "\n",
    "    bands.append(\"ndvi\")\n",
    "    bands.append(\"ndmi\")\n",
    "\n",
    "    print(\"\\nall bands are scaling to 8-bit images...\\n\")\n",
    "    band_names_scaled = []\n",
    "    for band_name in band_names:\n",
    "        print(band_name)\n",
    "        scaled_name = scale_img(f\"{band_name}.tif\")\n",
    "        band_names_scaled.append(scaled_name)\n",
    "\n",
    "    print(\"\\nall bands are being merged...\\n\")\n",
    "    print(band_names_scaled)\n",
    "\n",
    "    merge(save_file_merged, *band_names_scaled)\n",
    "\n",
    "    for item in os.listdir(img_folder):\n",
    "        if item.endswith(\".tif\"):\n",
    "            os.remove(join(img_folder, item))\n",
    "    return save_file_merged\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19f7b662",
   "metadata": {},
   "outputs": [],
   "source": [
    "tiles_prefix = \"\".join(tiles)\n",
    "tiles_prefix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f418c30",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "os.makedirs(PREPARED_DATA_FOLDER, exist_ok=True)\n",
    "\n",
    "\n",
    "start_merged = []\n",
    "end_merged = []\n",
    "\n",
    "for start_date_folder, end_date_folder in zip(start_date_folders, end_date_folders):\n",
    "    \n",
    "    start_date_merged_path = prepare_data(start_date_folder, PREPARED_DATA_FOLDER)\n",
    "    end_date_merged_path = prepare_data(end_date_folder, PREPARED_DATA_FOLDER)\n",
    "    \n",
    "    start_merged.append(start_date_merged_path)\n",
    "    end_merged.append(end_date_merged_path)\n",
    "\n",
    "start_date_merged_path = stitch_tiles(start_merged, out_raster_path=os.path.join(PREPARED_DATA_FOLDER, f'all_merged_{dates[0]}_date_{tiles_prefix}.tif'))\n",
    "end_date_merged_path = stitch_tiles(end_merged, out_raster_path=os.path.join(PREPARED_DATA_FOLDER, f'all_merged_{dates[1]}_date_{tiles_prefix}.tif'))\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25bcd915",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "start_date_crop_path = crop_aoi_from_raster(\n",
    "                                    raster_path=start_date_merged_path,\n",
    "                                    aoi_gdf=aoi,\n",
    "                                    output_path=os.path.join(PREPARED_DATA_FOLDER,f'prepared_{dates[0]}_date_{tiles_prefix}_crop.tif')\n",
    "                                    )\n",
    "end_date_crop_path = crop_aoi_from_raster(\n",
    "                    raster_path=end_date_merged_path,\n",
    "                    aoi_gdf=aoi,\n",
    "                    output_path=os.path.join(PREPARED_DATA_FOLDER,f'prepared_{dates[1]}_date_{tiles_prefix}_crop.tif')\n",
    "                    )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6fcc0194",
   "metadata": {},
   "source": [
    "### 3.1 Prepare clouds tif files for postprocessing step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10c6c088",
   "metadata": {},
   "outputs": [],
   "source": [
    "from time_dependent.data_prepare.prepare_clouds import (detect_clouds, merge,\n",
    "                                                        search_band, to_tiff)\n",
    "\n",
    "\n",
    "def prepare_clouds(data_folder, save_path):\n",
    "    folder_basename = os.path.basename(data_folder)\n",
    "    print(folder_basename)\n",
    "    bands, band_names = [\n",
    "        \"B01\",\n",
    "        \"B02\",\n",
    "        \"B04\",\n",
    "        \"B05\",\n",
    "        \"B08\",\n",
    "        \"B8A\",\n",
    "        \"B09\",\n",
    "        \"B10\",\n",
    "        \"B11\",\n",
    "        \"B12\",\n",
    "    ], []\n",
    "\n",
    "    for band in bands:\n",
    "        band_names.append(join(data_folder, search_band(band, data_folder, \"jp2\")))\n",
    "\n",
    "    print(\"\\nall bands are converting to *tif...\\n\")\n",
    "\n",
    "    for band_name in band_names:\n",
    "        print(band_name[-3:])\n",
    "        to_tiff(f\"{band_name}.jp2\")\n",
    "\n",
    "    print(\"\\n all bands are being merged...\\n\")\n",
    "\n",
    "    save_file_merged = join(save_path, f\"{folder_basename}_full_merged.tif\")\n",
    "    print(len(band_names))\n",
    "    merge(save_file_merged, *band_names)\n",
    "\n",
    "    save_file_clouds = join(save_path, f\"{folder_basename}_clouds.tiff\")\n",
    "    detect_clouds(save_file_merged, save_file_clouds)\n",
    "    os.remove(save_file_merged)\n",
    "\n",
    "    for item in os.listdir(data_folder):\n",
    "        if item.endswith(\".tif\"):\n",
    "            os.remove(join(data_folder, item))\n",
    "\n",
    "    # os.system(f'rm {join(granule_folder, tile_folder, 'IMG_DATA')}*.jp2')\n",
    "    print(\"\\ntemp files have been deleted\\n\")\n",
    "    return save_file_clouds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62292eb1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16c0ab62",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "os.makedirs(CLOUD_DATA_FOLDER, exist_ok=True)\n",
    "\n",
    "start_clouds = []\n",
    "end_clouds = []\n",
    "\n",
    "for start_date_folder, end_date_folder in zip(start_date_folders, end_date_folders):\n",
    "    start_date_clouds_path = prepare_clouds(start_date_folder, CLOUD_DATA_FOLDER) \n",
    "    end_date_clouds_path = prepare_clouds(end_date_folder, CLOUD_DATA_FOLDER)\n",
    "    \n",
    "    start_clouds.append(start_date_clouds_path)\n",
    "    end_clouds.append(end_date_clouds_path)\n",
    "    \n",
    "start_date_clouds_path = stitch_tiles(start_clouds, out_raster_path=os.path.join(CLOUD_DATA_FOLDER, f'all_clouds_{dates[0]}_date_{tiles_prefix}.tif'))\n",
    "end_date_clouds_path = stitch_tiles(end_clouds, out_raster_path=os.path.join(CLOUD_DATA_FOLDER, f'all_clouds_{dates[1]}_date_{tiles_prefix}.tif'))\n",
    "\n",
    "    \n",
    "for filename in os.listdir(CLOUD_DATA_FOLDER):\n",
    "    if not filename.startswith('all_'):\n",
    "        os.remove(os.path.join(CLOUD_DATA_FOLDER, filename))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3802eea",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_date_clouds_crop_path = crop_aoi_from_raster(\n",
    "                                    raster_path=start_date_clouds_path,\n",
    "                                    aoi_gdf=aoi,\n",
    "                                    output_path=os.path.join(CLOUD_DATA_FOLDER,f'clouds_{dates[0]}_date_{tiles_prefix}_crop.tif')\n",
    "                                    )\n",
    "end_date_clouds_crop_path = crop_aoi_from_raster(\n",
    "                    raster_path=end_date_clouds_path,\n",
    "                    aoi_gdf=aoi,\n",
    "                    output_path=os.path.join(CLOUD_DATA_FOLDER,f'clouds_{dates[0]}_date_{tiles_prefix}_crop.tif')\n",
    "                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2e484b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "clouds_files = [start_date_clouds_crop_path, end_date_clouds_crop_path]\n",
    "clouds_files"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "dd20af44",
   "metadata": {},
   "source": [
    "### 3.2 Preparing forest landcover data, also needed for postprocessing stage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2f9247c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LandcoverPolygons:\n",
    "    \"\"\"\n",
    "    LandcoverPolygon class to access forest polygons. Before usage,\n",
    "    be sure that SENTINEL_TILES file is downloaded.\n",
    "    SENTINEL_TILES_POLYGONS = 'https://sentinel.esa.int/documents/247904/1955685/S2A_OPER_GIP_TILPAR_MPC__20151209T095117_V20150622T000000_21000101T000000_B00.kml'\n",
    "\n",
    "    :param tile: tile name (str), e.g. '36UYA'\n",
    "    :param crs: coordinate system (str), e.g. 'EPSG:4326'\n",
    "\n",
    "    :return polygons: list of forest polygons within a tile in CRS of a S2A image\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, tiles, crs, year, aoi):\n",
    "        self.tiles = tiles\n",
    "        self.crs = crs\n",
    "\n",
    "        landcover_tiles = set([tile[:3] for tile in self.tiles])\n",
    "        self.LANDCOVER_GEOJSON = prepare_landcover(\n",
    "            year, landcover_tiles, LANDCOVER_POLYGONS_PATH, aoi\n",
    "        )\n",
    "        gpd.io.file.fiona.drvsupport.supported_drivers[\"KML\"] = \"rw\"\n",
    "\n",
    "    def get_polygon(self):\n",
    "        polygon_path = os.path.join(LANDCOVER_POLYGONS_PATH, f\"{''.join(self.tiles)}.geojson\")\n",
    "        logging.info(f\"LANDCOVER_POLYGONS_PATH: {polygon_path}\")\n",
    "        if os.path.isfile(polygon_path):\n",
    "            logging.info(f\"{self.tiles} forests polygons file exists.\")\n",
    "            polygons = gpd.read_file(polygon_path)\n",
    "        else:\n",
    "            logging.info(\n",
    "                f\"{self.tiles} forests polygons file does not exist. Creating polygons...\"\n",
    "            )\n",
    "            polygons = self.create_polygon()\n",
    "\n",
    "        if len(polygons) > 0:\n",
    "            polygons = polygons.to_crs(self.crs)\n",
    "            polygons = list(polygons[\"geometry\"])\n",
    "        else:\n",
    "            logging.info(\"No forests polygons.\")\n",
    "        return polygons\n",
    "\n",
    "    def create_polygon(self):\n",
    "        polygons = []\n",
    "        if os.path.isfile(SENTINEL_TILES):\n",
    "            logging.info(\n",
    "                f\"read forests_polygons_file: {SENTINEL_TILES}, for tile {self.tiles}\"\n",
    "            )\n",
    "\n",
    "            sentinel_tiles = gpd.read_file(SENTINEL_TILES, driver=\"KML\")\n",
    "            sentinel_tiles = sentinel_tiles[sentinel_tiles[\"Name\"].isin(self.tiles)]\n",
    "\n",
    "            logging.info(f\"sentinel_tiles for {self.tiles}: {sentinel_tiles}\")\n",
    "\n",
    "            bounding_polygon = sentinel_tiles[\"geometry\"].values[0]\n",
    "            polygons = gpd.read_file(self.LANDCOVER_GEOJSON)\n",
    "            polygons = polygons[polygons[\"geometry\"].intersects(bounding_polygon)]\n",
    "            polygon_path = os.path.join(LANDCOVER_POLYGONS_PATH, f\"{''.join(self.tiles)}.geojson\")\n",
    "\n",
    "            logging.info(f\"forests_polygons_file_path: {polygon_path}\")\n",
    "\n",
    "            if polygons.empty:\n",
    "                return polygons\n",
    "            polygons.to_file(polygon_path, driver=\"GeoJSON\")\n",
    "        else:\n",
    "            logging.error(f\"{SENTINEL_TILES} doth not exists\")\n",
    "            raise FileNotFoundError\n",
    "        return polygons\n",
    "    \n",
    "\n",
    "def download_landcover():    \n",
    "    LANDCOVER_OUTPUT_DIR = os.path.join(SATELLITE_CACHE_FOLDER, \"landcover_result\")\n",
    "    os.system(\"AOI='{0}' START_DATE='{1}' END_DATE='{2}' SENTINEL_CACHE='{3}' OUTPUT_FOLDER='{4}' jupyter nbconvert --inplace --to=notebook --execute ./sip_landcover/landcover.ipynb\".format(AOI, START_DATE, END_DATE, SATELLITE_CACHE_FOLDER, LANDCOVER_OUTPUT_DIR))\n",
    "    \n",
    "    output_path = os.path.join(LANDCOVER_OUTPUT_DIR, f\"{START_DATE}_{END_DATE}.tif\")\n",
    "    return output_path\n",
    "    \n",
    "def prepare_landcover(year, landcover_tiles, output_path, aoi):\n",
    "    filename = f\"{''.join(landcover_tiles)}_{year}_landcover\"\n",
    "\n",
    "    for file in os.listdir(output_path):\n",
    "        if file != LANDCOVER_FILENAME and not file.endswith(\"_landcover.geojson\"):\n",
    "            os.remove(os.path.join(output_path, file))\n",
    "\n",
    "#     for file in os.listdir(output_path):\n",
    "#         if filename in file:\n",
    "#             return os.path.join(output_path, file)\n",
    "    try:\n",
    "        landcover_annual(year, landcover_tiles, output_path, aoi)\n",
    "        landcover_geojson = create_landcover_gdf(output_path, output_path, filename)\n",
    "    except Exception:\n",
    "        raise Exception('Cannot create landcover gdf')\n",
    "    return landcover_geojson\n",
    "\n",
    "\n",
    "def landcover_annual(year, landcover_tiles, output_path, aoi):\n",
    "    # landcover_classes = {\n",
    "    #    1: \"Water\",\n",
    "    #    2: \"Trees\",\n",
    "    #    4: \"Flooded vegetation\",\n",
    "    #    5: \"Crops\",\n",
    "    #    7: \"Built Area\",\n",
    "    #    8: \"Bare ground\",\n",
    "    #   9: \"Snow/Ice\",\n",
    "    #   10: \"Clouds\",\n",
    "    #    11: \"Rangeland\"\n",
    "    # }\n",
    "    \n",
    "    colors_classes = {\n",
    "        \"65,155,223\": 1,\n",
    "        \"57,125,73\": 2,\n",
    "        \"122,135,198\": 4,\n",
    "        \"228,150,53\": 5,\n",
    "        \"196,50,27\": 7,\n",
    "        \"165,155,143\": 8,\n",
    "        \"168,235,255\": 9,\n",
    "        \"97,97,97\": 10,\n",
    "        \"227,226,195\": 11\n",
    "    }\n",
    "    \n",
    "    EPSG = \"EPSG:4326\"\n",
    "    landcover_downloaded_path = download_landcover()\n",
    "    \n",
    "    with rasterio.open(landcover_downloaded_path) as src:        \n",
    "        landcover_colored = src.read()\n",
    "        profile = src.profile\n",
    "        \n",
    "        landcover_crop = np.zeros((1, *landcover_colored.shape[-2:]))\n",
    "        \n",
    "        for color, landcover_class in colors_classes.items():\n",
    "            rgb_array = np.array(color.split(','), dtype=np.uint8)\n",
    "            class_indexes = np.where(np.all(landcover_colored==rgb_array[:, None, None], axis=0))\n",
    "\n",
    "            landcover_crop[:, class_indexes[0], class_indexes[1]] = landcover_class\n",
    "        \n",
    "        profile['count'] = 1\n",
    "        \n",
    "        crop_name = os.path.join(\n",
    "                    output_path, \"landcover_crop.tif\"\n",
    "                )\n",
    "        \n",
    "        with rasterio.open(crop_name, \"w\", **profile, nbits=1) as dst:\n",
    "            dst.write(np.where(landcover_crop == 2, 1, 0).astype(np.uint8))\n",
    "        \n",
    "    crops = [crop_name]\n",
    "\n",
    "    landcover_name = os.path.join(output_path, f\"landcover{year}.tif\")\n",
    "    listToStr = \" \".join(crops)\n",
    "    os.system(\n",
    "        \" \".join(\n",
    "            [\n",
    "                f\"gdalwarp --config GDAL_CACHEMAX 3000 -wm 3000 -t_srs {EPSG}\",\n",
    "                listToStr,\n",
    "                landcover_name,\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "    \n",
    "    print(f\"{landcover_name} was merged\")\n",
    "\n",
    "\n",
    "def rescale(img, ratio):\n",
    "    width = int(img.shape[1] * ratio)\n",
    "    height = int(img.shape[0] * ratio)\n",
    "    dim = (width, height)\n",
    "    resized = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)\n",
    "    return resized\n",
    "\n",
    "\n",
    "def mask_to_polygons(mask: np.ndarray, transform) -> MultiPolygon:\n",
    "    \"\"\"\n",
    "    Converts raster mask to shapely MultiPolygon\n",
    "    \"\"\"\n",
    "    polygons = []\n",
    "    shapes = features.shapes(\n",
    "        mask.astype(np.uint8), mask=(mask > 0), transform=transform\n",
    "    )\n",
    "\n",
    "    for shape, _ in shapes:\n",
    "        polygons.append(shapely.geometry.shape(shape))\n",
    "\n",
    "    polygons = MultiPolygon(polygons)\n",
    "    if not polygons.is_valid:\n",
    "        polygons = polygons.buffer(0)\n",
    "        if polygons.type == \"Polygon\":\n",
    "            polygons = MultiPolygon([polygons])\n",
    "    return polygons\n",
    "\n",
    "\n",
    "def create_landcover_gdf(landcover_dir, output_dir, filename):\n",
    "    landcover_names = [\n",
    "        name for name in os.listdir(landcover_dir) if name.endswith(\"_crop.tif\")\n",
    "    ]\n",
    "    gdfs = []\n",
    "    for name in tqdm(landcover_names):\n",
    "        lc_fullpath = os.path.join(landcover_dir, name)\n",
    "        print(lc_fullpath)\n",
    "        with rasterio.open(lc_fullpath, \"r\") as src:\n",
    "            data = src.read().squeeze()\n",
    "            data = rescale(data, 0.5)\n",
    "            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (10, 10))\n",
    "            data = cv2.morphologyEx(data, cv2.MORPH_CLOSE, kernel)\n",
    "            data = cv2.morphologyEx(data, cv2.MORPH_OPEN, kernel)\n",
    "            data = rescale(data, 2)\n",
    "\n",
    "            current_polygons = list(mask_to_polygons(data, src.transform))\n",
    "            current_polygons = [poly for poly in current_polygons if poly.area]\n",
    "\n",
    "            current_areas = [poly.area for poly in current_polygons]\n",
    "            current_tilename = [name] * len(current_areas)\n",
    "\n",
    "            crs = src.crs\n",
    "            current_gdf = gpd.GeoDataFrame(\n",
    "                {\n",
    "                    \"area\": current_areas,\n",
    "                    \"names\": current_tilename,\n",
    "                    \"geometry\": current_polygons,\n",
    "                },\n",
    "                crs=crs,\n",
    "            )\n",
    "            current_gdf.to_crs(\"EPSG:4326\", inplace=True)\n",
    "            gdfs.append(current_gdf.copy())\n",
    "\n",
    "    gdf = gpd.GeoDataFrame(pd.concat(gdfs, ignore_index=True), crs=\"EPSG:4326\")\n",
    "\n",
    "    os.makedirs(output_dir, exist_ok=True)\n",
    "    output_path = os.path.join(output_dir, f\"{filename}.geojson\")\n",
    "    gdf.to_file(output_path, driver=\"GeoJSON\")\n",
    "    return output_path\n",
    "\n",
    "\n",
    "def weights_exists_or_download(path, file_id):\n",
    "    if not Path(path).exists():\n",
    "        creds_file = os.environ.get(\"CREDENTIAL_FILE\")\n",
    "        creds = service_account.Credentials.from_service_account_file(\n",
    "            creds_file, scopes=SCOPES\n",
    "        )\n",
    "\n",
    "        service = build(\"drive\", \"v3\", credentials=creds)\n",
    "        request = service.files().get_media(fileId=file_id)\n",
    "\n",
    "        fh = io.FileIO(\"unet_v4.pth\", mode=\"wb\")\n",
    "        downloader = MediaIoBaseDownload(fh, request)\n",
    "        done = False\n",
    "        while done is False:\n",
    "            status, done = downloader.next_chunk()\n",
    "            print(f\"Download {int(status.progress() * 100)}\")\n",
    "\n",
    "    return path"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "84cfe616",
   "metadata": {},
   "source": [
    "### 4. Unet-diff inference on the prepared data and then postprocessing on cloud and forest landcover data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41fa8eb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n",
    "os.environ.get(\"CUDA_VISIBLE_DEVICES\", \"0\")\n",
    "\n",
    "logging.basicConfig(format=\"%(asctime)s %(message)s\")\n",
    "\n",
    "\n",
    "def predict(image_tensor, model, channels, neighbours, size, device):\n",
    "    image_shape = 1, count_channels(channels) * neighbours, size, size\n",
    "    prediction, _ = model.predict(\n",
    "        image_tensor.view(image_shape).to(device, dtype=torch.float)\n",
    "    )\n",
    "    result = prediction.view(size, size).detach().cpu().numpy()\n",
    "    return result\n",
    "\n",
    "\n",
    "def diff(img1, img2):\n",
    "    img2 = match_histograms(img2, img1, multichannel=True)\n",
    "    difference = (img1 - img2) / (img1 + img2)\n",
    "    difference = (difference + 1) * 127\n",
    "    return np.concatenate(\n",
    "        (difference.astype(np.uint8), img1.astype(np.uint8), img2.astype(np.uint8)),\n",
    "        axis=-1,\n",
    "    )\n",
    "\n",
    "\n",
    "def mask_postprocess(mask):\n",
    "    kernel = np.ones((3, 3), np.uint8)\n",
    "    erosion = cv2.erode(mask, kernel, iterations=1)\n",
    "    kernel = np.ones((5, 5), np.uint8)\n",
    "    closing = cv2.morphologyEx(erosion, cv2.MORPH_CLOSE, kernel)\n",
    "    return closing\n",
    "\n",
    "\n",
    "def predict_raster(\n",
    "    img_current,\n",
    "    img_previous,\n",
    "    channels,\n",
    "    network=\"unet-diff\",\n",
    "    model_weights_path=MODEL_PATH,\n",
    "    input_size=56,\n",
    "    neighbours=3,\n",
    "):\n",
    "    model, device = load_model(network, model_weights_path)\n",
    "\n",
    "    with rasterio.open(img_current) as source_current, rasterio.open(\n",
    "        img_previous\n",
    "    ) as source_previous:\n",
    "        meta = source_current.meta\n",
    "        meta[\"count\"] = 1\n",
    "        \n",
    "        height = (source_current.height \n",
    "                  if  source_current.height % input_size == 0 \n",
    "                  else (source_current.height + (input_size - source_current.height % input_size)))\n",
    "        width = (source_current.width \n",
    "                  if  source_current.width % input_size == 0 \n",
    "                  else (source_current.width + (input_size - source_current.width % input_size)))\n",
    "        clearcut_mask = np.zeros((height, width))\n",
    "        for i in tqdm(range(width // input_size)):\n",
    "            for j in range(height // input_size):\n",
    "                bottom_row = i * input_size\n",
    "                upper_row = (i + 1) * input_size\n",
    "                left_column = j * input_size\n",
    "                right_column = (j + 1) * input_size\n",
    "\n",
    "                corners = [\n",
    "                    source_current.xy(bottom_row, left_column),\n",
    "                    source_current.xy(bottom_row, right_column),\n",
    "                    source_current.xy(upper_row, right_column),\n",
    "                    source_current.xy(upper_row, left_column),\n",
    "                    source_current.xy(bottom_row, left_column),\n",
    "                ]\n",
    "\n",
    "                window = Window(bottom_row, left_column, input_size, input_size)\n",
    "                image_current = reshape_as_image(source_current.read(window=window))\n",
    "                image_previous = reshape_as_image(source_previous.read(window=window))\n",
    "                              \n",
    "                h, w = image_current.shape[0:2]\n",
    "                image_current = np.pad(image_current,\n",
    "                                       pad_width=((0,input_size-h),\n",
    "                                                  (0,input_size-w),\n",
    "                                                  (0,0)))\n",
    "                image_previous = np.pad(image_previous,\n",
    "                                       pad_width=((0,input_size-h),\n",
    "                                                  (0,input_size-w),\n",
    "                                                  (0,0)))\n",
    "\n",
    "                if np.any(image_current) and np.any(image_previous):\n",
    "                    difference_image = diff(image_current, image_previous)\n",
    "                    image_tensor = transforms.ToTensor()(\n",
    "                        difference_image.astype(np.uint8)\n",
    "                    ).to(device, dtype=torch.float)\n",
    "\n",
    "                    predicted = predict(\n",
    "                        image_tensor, model, channels, neighbours, input_size, device\n",
    "                    )\n",
    "                    predicted = mask_postprocess(predicted)\n",
    "                    clearcut_mask[\n",
    "                        left_column:right_column, bottom_row:upper_row\n",
    "                    ] += predicted\n",
    "                    \n",
    "    meta[\"dtype\"] = \"float32\"\n",
    "    return clearcut_mask.astype(np.float32), meta\n",
    "\n",
    "\n",
    "def count_channels(channels):\n",
    "    count = 0\n",
    "    for ch in channels:\n",
    "        ch = ch.lower()\n",
    "        if ch == \"rgb\":\n",
    "            count += 3\n",
    "        elif ch in [\"ndvi\", \"ndmi\", \"b08\", \"b8a\", \"b11\", \"b12\"]:\n",
    "            count += 1\n",
    "        else:\n",
    "            raise Exception(\"{} channel is unknown!\".format(ch))\n",
    "\n",
    "    return count\n",
    "\n",
    "\n",
    "def scale(tensor, max_value):\n",
    "    max_ = tensor.max()\n",
    "    if max_ > 0:\n",
    "        return tensor / max_ * max_value\n",
    "    return tensor\n",
    "\n",
    "\n",
    "def save_raster(raster_array, meta, save_path, filename):\n",
    "    if not os.path.exists(save_path):\n",
    "        os.makedirs(save_path, exist_ok=True)\n",
    "        logging.info(\"Data directory created.\")\n",
    "\n",
    "    save_path = os.path.join(save_path, filename)\n",
    "\n",
    "    with rasterio.open(f\"{save_path}.tif\", \"w\", **meta) as dst:\n",
    "        for i in range(1, meta[\"count\"] + 1):\n",
    "            dst.write(raster_array, i)\n",
    "\n",
    "\n",
    "\n",
    "def polygonize(raster_array, meta, transform=True, mode=cv2.RETR_TREE):\n",
    "    raster_array = (raster_array * 255).astype(np.uint8)\n",
    "\n",
    "    contours, _ = cv2.findContours(raster_array, mode, cv2.CHAIN_APPROX_SIMPLE)\n",
    "\n",
    "    polygons = []\n",
    "    for i in tqdm(range(len(contours))):\n",
    "        c = contours[i]\n",
    "        n_s = (c.shape[0], c.shape[2])\n",
    "        if n_s[0] > 2:\n",
    "            if transform:\n",
    "                polys = [tuple(i) * meta[\"transform\"] for i in c.reshape(n_s)]\n",
    "            else:\n",
    "                polys = [tuple(i) for i in c.reshape(n_s)]\n",
    "            polygons.append(Polygon(polys))\n",
    "\n",
    "    return polygons\n",
    "\n",
    "def dump_no_data_geosjon(polygon, geojson_path):\n",
    "    feature = geojson.Feature(geometry=polygon, properties=dict(label='No data', style=dict(color='red')))\n",
    "    feature['start_date'] = START_DATE\n",
    "    feature['end_date'] = END_DATE\n",
    "    feature['name'] = 'No deforestation found\\nNo data available'\n",
    "\n",
    "    with open(geojson_path, 'w') as f:\n",
    "        geojson.dump(feature, f)\n",
    "        \n",
    "\n",
    "def save_polygons(polygons, \n",
    "                  save_path, \n",
    "                  filename, \n",
    "                  aoi):\n",
    "    \n",
    "    if not os.path.exists(save_path):\n",
    "        os.makedirs(save_path, exist_ok=True)\n",
    "        logging.info(\"Data directory created.\")\n",
    "        \n",
    "    geojson_output_path = os.path.join(save_path, f\"{filename}.geojson\")\n",
    "    if len(polygons) == 0:\n",
    "        logging.info(\"no_polygons detected\")\n",
    "        dump_no_data_geosjon(aoi.geometry[0], geojson_output_path)\n",
    "        \n",
    "    else:\n",
    "        logging.info(f\"{filename} saved.\")\n",
    "        print(f\"{filename} saved.\")\n",
    "        polygons['start_date'] = START_DATE\n",
    "        polygons['end_date'] = END_DATE\n",
    "        \n",
    "        polygons = polygons.to_crs(\"EPSG:4326\")\n",
    "        \n",
    "        style = {\n",
    "            \"color\": \"#e80e27\",\n",
    "            \"stroke\": \"#e80e27\",\n",
    "            \"stroke-width\": 2\n",
    "        }\n",
    "\n",
    "        # Assign the style to each polygon using apply()\n",
    "        polygons[\"style\"] = polygons.apply(lambda row: style, axis=1)\n",
    "        \n",
    "        polygons.to_file(geojson_output_path, driver=\"GeoJSON\")\n",
    "\n",
    "\n",
    "def intersection_poly(test_poly, mask_poly):\n",
    "    intersecion_score = False\n",
    "    if test_poly.is_valid and mask_poly.is_valid:\n",
    "        intersection_result = test_poly.intersection(mask_poly)\n",
    "        if not intersection_result.is_valid:\n",
    "            intersection_result = intersection_result.buffer(0)\n",
    "        if not intersection_result.is_empty:\n",
    "            intersecion_score = True\n",
    "    return intersecion_score\n",
    "\n",
    "\n",
    "def morphological_transform(img):\n",
    "    kernel = np.ones((5, 5), np.uint8)\n",
    "    closing = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)\n",
    "\n",
    "    kernel = np.ones((3, 3), np.uint8)\n",
    "    closing = cv2.dilate(closing, kernel, iterations=1)\n",
    "    return closing\n",
    "\n",
    "\n",
    "def postprocessing(tiles, cloud_files, clearcuts, src_crs, year, aoi):\n",
    "    def get_intersected_polygons(polygons, masks, mask_column_name):\n",
    "        \"\"\"Finding in GeoDataFrame with clearcuts the masked polygons.\n",
    "\n",
    "        :param polygons: GeoDataFrame with clearcuts and mask columns\n",
    "        :param masks: list of masks (e.g., polygons of clouds)\n",
    "        :param mask_column_name: name of mask column in polygons GeoDataFrame\n",
    "\n",
    "        :return: GeoDataFrame with filled mask flags in corresponding column\n",
    "        \"\"\"\n",
    "        masked_values = []\n",
    "        if len(masks) > 0:\n",
    "            centroids = [[mask.centroid.x, mask.centroid.y] for mask in masks]\n",
    "            kdtree = spatial.KDTree(centroids)\n",
    "            for _, clearcut in polygons.iterrows():\n",
    "                polygon = clearcut[\"geometry\"]\n",
    "                _, idxs = kdtree.query(polygon.centroid, k=NEAREST_POLYGONS_NUMBER)\n",
    "                masked_value = 0\n",
    "                for idx in idxs:\n",
    "                    if idx >= len(masks):\n",
    "                        break\n",
    "                    if intersection_poly(polygon, masks[idx].buffer(0)):\n",
    "                        masked_value = 1\n",
    "                        break\n",
    "                masked_values.append(masked_value)\n",
    "        polygons[mask_column_name] = masked_values\n",
    "        return polygons\n",
    "\n",
    "    landcover = LandcoverPolygons(tiles, src_crs, year, aoi)\n",
    "    forest_polygons = landcover.get_polygon()\n",
    "\n",
    "    #     cloud_files = [f\"{img_path}/{tile}_{i}/clouds.tiff\" for i in range(DATES_FOR_TILE)]\n",
    "    cloud_polygons = []\n",
    "    for cloud_file in cloud_files:\n",
    "        with rasterio.open(cloud_file) as src:\n",
    "            clouds = src.read(1)\n",
    "            meta = src.meta\n",
    "        clouds = morphological_transform(clouds)\n",
    "        clouds = (clouds > CLOUDS_PROBABILITY_THRESHOLD).astype(np.uint8)\n",
    "        if clouds.sum() > 0:\n",
    "            cloud_polygons.extend(polygonize(clouds, meta, mode=cv2.RETR_LIST))\n",
    "\n",
    "    n_clearcuts = len(clearcuts)\n",
    "    polygons = {\n",
    "        \"geometry\": gpd.GeoSeries(clearcuts),\n",
    "        \"forest\": np.zeros(n_clearcuts),\n",
    "        \"clouds\": np.zeros(n_clearcuts),\n",
    "    }\n",
    "\n",
    "    polygons = geopandas.GeoDataFrame(polygons, crs=src_crs)\n",
    "\n",
    "    if len(cloud_polygons) > 0:\n",
    "        polygons = get_intersected_polygons(polygons, cloud_polygons, \"clouds\")\n",
    "    else:\n",
    "        print(\"Clouds with specified CLOUDS_PROBABILITY_THRESHOLD not found\")\n",
    "    polygons = get_intersected_polygons(polygons, forest_polygons, \"forest\")\n",
    "    polygons = polygons.loc[(polygons.forest == 1) & (polygons.clouds == 0)]\n",
    "    \n",
    "    return polygons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56fd33bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def deforestation_inference(\n",
    "    img_start_path,\n",
    "    img_end_path,\n",
    "    tiles,\n",
    "    network=\"unet-diff\",\n",
    "    model_weights_path=MODEL_PATH,\n",
    "    save_path=OUTPUT_FOLDER,\n",
    "    channels=[\"RGB\", \"B08\", \"B8A\", \"B11\", \"B12\", \"NDVI\", \"NDMI\"],\n",
    "    threshold=0.4,\n",
    "    polygonize_only=False,\n",
    "    clouds_files=clouds_files,\n",
    "    landcover_aoi=aoi,\n",
    "):\n",
    "    predicted_filename = f\"predicted_{''.join(dates)}_{tiles_prefix}\"\n",
    "\n",
    "    raster_array, meta = predict_raster(\n",
    "        img_end_path, img_start_path, channels, network, model_weights_path\n",
    "    )\n",
    "\n",
    "    logging.info(\"Polygonize raster array of clearcuts...\")\n",
    "    clearcuts = polygonize(raster_array > threshold, meta)\n",
    "    logging.info(\"Filter polygons of clearcuts\")\n",
    "    polygons = postprocessing(\n",
    "        tiles,\n",
    "        clouds_files,\n",
    "        clearcuts,\n",
    "        meta[\"crs\"],\n",
    "        year,\n",
    "        landcover_aoi,\n",
    "    )\n",
    "    \n",
    "    save_polygons(polygons, \n",
    "                  save_path, \n",
    "                  predicted_filename, \n",
    "                  landcover_aoi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "090007ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "deforestation_inference(start_date_crop_path, end_date_crop_path, tiles, \n",
    "                        clouds_files=clouds_files,\n",
    "                        landcover_aoi=aoi)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5553579f",
   "metadata": {},
   "source": [
    "#### Delete all intermediate files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19beffde",
   "metadata": {},
   "outputs": [],
   "source": [
    "shutil.rmtree(SATELLITE_CACHE_FOLDER)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}