todo-fuse-postprocess.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1848ecc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import rasterio\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "import math\n",
    "import maxflow as mf\n",
    "import dask\n",
    "import scipy.ndimage\n",
    "import skimage.morphology\n",
    "from concurrent.futures import ThreadPoolExecutor\n",
    "\n",
    "#import settings.by_dop80c_1312.detectree_r1_deepforest_r1.setting as setting\n",
    "import settings.opendata_luftbild_dop60_1312.detectree_r1_deepforest_r1.setting as setting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1736ae9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "CONTINUE_MODE = True  # ignore already processed images\n",
    "SAVE_PROBA = False\n",
    "SAVE_IMAGE = True\n",
    "\n",
    "DETECTREE_PROBA_RESULT_PATH = setting.POSTPROCESS_SRC_DETECTREE_PREDICT_RESULT_PATH.joinpath('p')\n",
    "DEEPFOREST_BBOX_RESULT_PATH = setting.POSTPROCESS_SRC_DEEPFOREST_PREDICT_RESULT_PATH.joinpath('b')\n",
    "\n",
    "OUTPUT_PROBA_RESULT_PATH = setting.DATASET_PREDICT_RESULT_PATH.joinpath('p')\n",
    "OUTPUT_IMAGE_RESULT_PATH = setting.DATASET_PREDICT_RESULT_PATH.joinpath('r')\n",
    "\n",
    "MODEL_PARAMS_CLASSIFIER = setting.MODEL_PARAMS_CLASSIFIER\n",
    "\n",
    "FUSE_CONFIG = setting.POSTPROCESS_CONFIG\n",
    "\n",
    "CONCURRENCY = 4\n",
    "\n",
    "os.makedirs(OUTPUT_PROBA_RESULT_PATH, exist_ok=True)\n",
    "os.makedirs(OUTPUT_IMAGE_RESULT_PATH, exist_ok=True)\n",
    "\n",
    "print(\"detectree:\", DETECTREE_PROBA_RESULT_PATH)\n",
    "print(\"deepforest:\", DEEPFOREST_BBOX_RESULT_PATH)\n",
    "print(\"output:\", setting.DATASET_PREDICT_RESULT_PATH)\n",
    "print(\"fuse:\", FUSE_CONFIG)\n",
    "print(\"clf:\", MODEL_PARAMS_CLASSIFIER)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ad5d450",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dilation -> closing -> fill holes -> erosion\n",
    "def refine_morph(image, morph_size, tree_val, nontree_val):\n",
    "  labels = image[0,:,:,] == tree_val # turn to boolean type ndarray\n",
    "  refined = scipy.ndimage.morphology.binary_dilation(labels, \n",
    "                    structure=skimage.morphology.disk(morph_size))\n",
    "  refined = scipy.ndimage.morphology.binary_closing(refined, \n",
    "                    structure=skimage.morphology.disk(morph_size*2))\n",
    "  refined = scipy.ndimage.morphology.binary_fill_holes(refined, \n",
    "                    structure=skimage.morphology.disk(int(morph_size/2)))\n",
    "  refined = scipy.ndimage.morphology.binary_erosion(refined, \n",
    "                    structure=skimage.morphology.disk(morph_size))\n",
    "  return np.where(refined == True, tree_val, nontree_val).reshape(1, image.shape[1], image.shape[2])\n",
    "\n",
    "# max-flow min0cut refinement\n",
    "def refine_mfmc(img_shape, p_tree, refine_int_rescale, refine_beta, tree_val, nontree_val):\n",
    "  g = mf.Graph[int]()\n",
    "  node_ids = g.add_grid_nodes(img_shape)\n",
    "  D_tree = (refine_int_rescale * np.log(1.0-p_tree)).astype(int)\n",
    "  D_nontree = (refine_int_rescale * np.log(p_tree)).astype(int)\n",
    "  MOORE_NEIGHBORHOOD_ARR = np.array([[0, 0, 0], [0, 0, 1], [1, 1, 1]])\n",
    "  g.add_grid_edges(node_ids, refine_beta,\n",
    "                    structure=MOORE_NEIGHBORHOOD_ARR)\n",
    "  g.add_grid_tedges(node_ids, D_tree, D_nontree)\n",
    "  g.maxflow()\n",
    "  refined = np.full(img_shape, nontree_val)\n",
    "  refined[g.get_grid_segments(node_ids)] = tree_val\n",
    "  return refined\n",
    "  \n",
    "def fuse(src_proba_path, src_bbox_path, output_proba_path, output_img_path, \n",
    "         radius_min_thres, radius_max_thres, score_min_thres, morph_size,\n",
    "         bias_tree_factor, bias_nontree_factor):\n",
    "  save_proba = SAVE_PROBA\n",
    "  save_img = SAVE_IMAGE\n",
    "\n",
    "  if save_proba and output_proba_path.exists() and CONTINUE_MODE:\n",
    "    save_proba = False\n",
    "  if save_img and output_img_path.exists() and CONTINUE_MODE:\n",
    "    save_img = False\n",
    "\n",
    "  if not save_proba and not save_img:\n",
    "    print(\"PASS:\", src_proba_path.stem)\n",
    "    return\n",
    "\n",
    "  # load bbox calculated from deepforest\n",
    "  bbox_df = pd.read_pickle(src_bbox_path).sort_values(by='score', ascending=True) if src_bbox_path.exists() else None\n",
    "\n",
    "  # load image of tree pixel probability from detectree\n",
    "  img_proba, img_transform, img_crs  = None, None, None\n",
    "  with rasterio.open(src_proba_path) as src:\n",
    "    img_transform = src.transform\n",
    "    img_crs = src.crs\n",
    "    img_proba = src.read()\n",
    "\n",
    "  # approximate an image of tree pixel probability according to bbox\n",
    "  if bbox_df is not None:\n",
    "    mask_score = np.zeros(img_proba.shape, dtype=img_proba.dtype)\n",
    "    for _, box in bbox_df.iterrows():\n",
    "      xmin, xmax, ymin, ymax = box.xmin, box.xmax, box.ymin, box.ymax\n",
    "      score = box.score\n",
    "      center_r, center_c = int((ymin + ymax)/2), int((xmin + xmax)/2)\n",
    "      r = math.ceil(min(ymax-ymin, xmax-xmin) / 2)\n",
    "      if r > radius_min_thres and r < radius_max_thres:\n",
    "        y, x = np.ogrid[-center_r:img_proba.shape[1]-center_r, -center_c:img_proba.shape[2]-center_c]\n",
    "        mask = x*x + y*y <= r*r\n",
    "        mask_score[:,mask] = score\n",
    "    # fuse probability and score to get a smooth coverage\n",
    "    mask_score -= score_min_thres\n",
    "    img_proba = np.where(mask_score >= 0, \n",
    "                        img_proba + (img_proba.max()-img_proba) * mask_score * bias_tree_factor, \n",
    "                        img_proba + (img_proba-img_proba.min()) * mask_score * bias_nontree_factor)\n",
    "    np.clip(img_proba, 0.0, 1.0, out=img_proba)\n",
    "  if save_proba: # float 1 channel image\n",
    "    with rasterio.open(output_proba_path, 'w', driver='GTiff',\n",
    "                  width=img_proba.shape[2], height=img_proba.shape[1],\n",
    "                  count=1, dtype=img_proba.dtype, nodata=0,\n",
    "                  transform=img_transform, crs=img_crs) as dst:\n",
    "      dst.write(img_proba.astype(img_proba.dtype))\n",
    "\n",
    "  if save_img: # uint8 1 channel image\n",
    "    #img_refined = refine_naive(img_proba, score_min_thres, MODEL_PARAMS_CLASSIFIER['tree_val'], MODEL_PARAMS_CLASSIFIER['nontree_val'])\n",
    "    img_refined = refine_mfmc(img_shape=img_proba.shape, p_tree=img_proba, \n",
    "                      refine_int_rescale=MODEL_PARAMS_CLASSIFIER['refine_int_rescale'],\n",
    "                      refine_beta=MODEL_PARAMS_CLASSIFIER['refine_beta'],\n",
    "                      tree_val=MODEL_PARAMS_CLASSIFIER['tree_val'], \n",
    "                      nontree_val=MODEL_PARAMS_CLASSIFIER['nontree_val'])\n",
    "    #img_refined = refine_morph(img_refined, morph_size=morph_size,\n",
    "    #                  tree_val=MODEL_PARAMS_CLASSIFIER['tree_val'], \n",
    "    #                  nontree_val=MODEL_PARAMS_CLASSIFIER['nontree_val'])\n",
    "    with rasterio.open(output_img_path, 'w', driver='GTiff',\n",
    "                    width=img_proba.shape[2], height=img_proba.shape[1],\n",
    "                    count=1, dtype=rasterio.uint8, nodata=0,\n",
    "                    transform=img_transform, crs=img_crs) as dst:\n",
    "      dst.write(img_refined.astype(np.uint8))\n",
    "  \n",
    "  print(\"PROCESSED:\", src_proba_path.stem)\n",
    "\n",
    "tasks = []\n",
    "for src_proba_path in DETECTREE_PROBA_RESULT_PATH.glob('*.tiff'):\n",
    "  src_bbox_path = DEEPFOREST_BBOX_RESULT_PATH.joinpath(src_proba_path.stem + \".pkl\")\n",
    "  output_proba_path = OUTPUT_PROBA_RESULT_PATH.joinpath(src_proba_path.stem + \".tiff\")\n",
    "  output_img_path = OUTPUT_IMAGE_RESULT_PATH.joinpath(src_proba_path.stem + \".tiff\")\n",
    "  delayed = dask.delayed(fuse)(src_proba_path, src_bbox_path, output_proba_path, output_img_path, **FUSE_CONFIG)\n",
    "  tasks.append(delayed)\n",
    "with dask.config.set(pool=ThreadPoolExecutor(CONCURRENCY)):  \n",
    "  dask.compute(*tasks)\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "d17c320a108ed353692a9d869777d22c4cfc8e217512b3979b490c9a0e17a049"
  },
  "kernelspec": {
   "display_name": "Python 3.9.6 64-bit",
   "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}