{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [Angle closure Glaucoma Evaluation Challenge](https://age.grand-challenge.org/Details/)\n",
    "## Scleral spur localization Baseline （FCN)\n",
    "\n",
    "- To keep model training stable, images with coordinate == -1, were removed.\n",
    "\n",
    "- For real inference, you MIGHT keep all images in val_file_path file."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "- Assume `Training100.zip` and `Validation_ASOCT_Image.zip` are stored @ `./AGE_challenge Baseline/datasets/`\n",
    "- Assume `weights` are stored @ `./AGE_challenge Baseline/weights/`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download ImageNet weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2019-08-06 14:16:30--  https://paddle-imagenet-models-name.bj.bcebos.com/ResNet34_pretrained.tar\n",
      "Resolving paddle-imagenet-models-name.bj.bcebos.com (paddle-imagenet-models-name.bj.bcebos.com)... 111.206.47.194, 202.106.5.21\n",
      "Connecting to paddle-imagenet-models-name.bj.bcebos.com (paddle-imagenet-models-name.bj.bcebos.com)|111.206.47.194|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 87470080 (83M) [application/x-tar]\n",
      "Saving to: ‘../weights/ResNet34_pretrained.tar’\n",
      "\n",
      "ResNet34_pretrained 100%[===================>]  83.42M  2.08MB/s    in 43s     \n",
      "\n",
      "2019-08-06 14:17:14 (1.93 MB/s) - ‘../weights/ResNet34_pretrained.tar’ saved [87470080/87470080]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# https://github.com/PaddlePaddle/models/tree/develop/PaddleCV/image_classification\n",
    "!rm ../weights/ResNet34_pretrained.tar \n",
    "!rm -rf ../weights/ResNet34_pretrained\n",
    "\n",
    "!wget -P ../weights/ https://paddle-imagenet-models-name.bj.bcebos.com/ResNet34_pretrained.tar \n",
    "!tar xvf ../weights/ResNet34_pretrained.tar -C ../weights/ > /dev/null # silent\n",
    "!rm ../weights/ResNet34_pretrained/fc*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Main Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, random, functools, math\n",
    "import cv2\n",
    "import numpy as np\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running Verify Fluid Program ... \n",
      "Your Paddle Fluid works well on SINGLE GPU or CPU.\n",
      "Your Paddle Fluid works well on MUTIPLE GPU or CPU.\n",
      "Your Paddle Fluid is installed successfully! Let's start deep Learning with Paddle Fluid now\n"
     ]
    }
   ],
   "source": [
    "import paddle\n",
    "import paddle.fluid as fluid\n",
    "import paddle.fluid.layers as FL\n",
    "import paddle.fluid.optimizer as FO\n",
    "fluid.install_check.run_check()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# preprocess: extract left/right label col in Training100_Location.xlsx\n",
    "# save to train_csv file\n",
    "data_root_path = \"../datasets/Training100/\"\n",
    "image_path = os.path.join(data_root_path, \"ASOCT_Image_loc\")\n",
    "\n",
    "train_file_path = os.path.join(data_root_path, \"loc_train_split.csv\")\n",
    "val_file_path = os.path.join(data_root_path, \"loc_val_split.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 32\n",
    "THREAD = 8\n",
    "BUF_SIZE = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove last global pooling and fullyconnect layer to enable FCN arch.\n",
    "# Standard ResNet Implement: \n",
    "# https://github.com/PaddlePaddle/models/blob/develop/PaddleCV/image_classification/models/resnet.py\n",
    "from resnet_modified import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Data Loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vflip_image(image):\n",
    "    return cv2.flip(image, flipCode=1)\n",
    "\n",
    "def gaussian_k(x0,y0, sigma, width, height):\n",
    "    \"\"\" Make a square gaussian kernel centered at (x0, y0) with sigma as SD.\n",
    "    \"\"\"\n",
    "    x = np.arange(0, width, 1, float) ## (width,)\n",
    "    y = np.arange(0, height, 1, float)[:, np.newaxis] ## (height,1)\n",
    "    return np.exp(-((x-x0)**2 + (y-y0)**2) / (2*sigma**2))\n",
    "\n",
    "def generate_hm(height, width, point, s=10):\n",
    "    \"\"\" Generate a full Heap Map for every landmarks in an array\n",
    "    Args:\n",
    "        height    : The height of Heat Map (the height of target output)\n",
    "        width     : The width  of Heat Map (the width of target output)\n",
    "        point    : (x,y)\n",
    "    \"\"\"\n",
    "    hm = gaussian_k(point[0], point[1], s, height, width)\n",
    "    return hm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reader(img_path, file_list, batch_size=32, shuffle=True, shuffle_seed=42):\n",
    "    def read_file_list():\n",
    "        batch_data = []\n",
    "        np.random.shuffle(file_list)\n",
    "        for line in file_list:\n",
    "            file_name, p_x, p_y = line.split(\",\")\n",
    "            batch_data.append([file_name, float(p_x), float(p_y)])\n",
    "            if len(batch_data) == batch_size:\n",
    "                yield batch_data\n",
    "                batch_data = []\n",
    "        if len(batch_data) != 0:\n",
    "            yield batch_data\n",
    "    return read_file_list\n",
    "\n",
    "def process_batch_data(input_data, mode, rotate=True, flip=True):\n",
    "    batch_data = []\n",
    "    for sample in input_data:\n",
    "        file, p_x, p_y = sample\n",
    "        \n",
    "        img = cv2.imread( file )\n",
    "        img = img[:, :, ::-1].astype('float32') / 255.0\n",
    "        \n",
    "        ratio = 256.0 / img.shape[0]\n",
    "        p_x, p_y = p_x * ratio, p_y * ratio\n",
    "        img = cv2.resize(img, (256, 256))\n",
    "\n",
    "        if mode == 'train':\n",
    "            img = img + np.random.randn(*img.shape) * 0.3 / 255 \n",
    "            if flip and np.random.randint(0,2):\n",
    "                img = vflip_image(img)\n",
    "                p_x = 256 - p_x\n",
    "        else:\n",
    "            pass\n",
    "        \n",
    "        hm = generate_hm(256, 256, (p_x, p_y))\n",
    "        img = img.transpose((2, 0, 1))\n",
    "        batch_data.append((img, hm))\n",
    "\n",
    "    return batch_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def data_loader(img_list, img_path, batch_size, order=False, mode='train'):\n",
    "    data_reader = reader(img_path, img_list, batch_size)\n",
    "    mapper = functools.partial(process_batch_data, mode=mode)\n",
    "    \n",
    "    data_reader = paddle.reader.shuffle(data_reader, 32)\n",
    "    \n",
    "    return paddle.reader.xmap_readers(\n",
    "        mapper, data_reader, THREAD, BUF_SIZE, order=order)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(train_file_path) as flist:\n",
    "    train_file_list = [os.path.join(image_path,line.strip()) for line in flist]\n",
    "\n",
    "with open(val_file_path) as flist:\n",
    "    val_file_list = [os.path.join(image_path,line.strip()) for line in flist]  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2586\n",
      "607\n",
      "../datasets/Training100/ASOCT_Image_loc/T0056-10_left.jpg,228.83365553922314,466.95960107867666\n"
     ]
    }
   ],
   "source": [
    "print(len(train_file_list))\n",
    "print(len(val_file_list))\n",
    "print(train_file_list[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.shuffle(train_file_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataloader = data_loader(train_file_list, image_path, BATCH_SIZE, False, mode='train')\n",
    "val_dataloader = data_loader(val_file_list, image_path, BATCH_SIZE, True, mode='val')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define model (compute graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def network():\n",
    "    data_shape = [3, 256, 256]\n",
    "    \n",
    "    model = ResNet34()\n",
    "    \n",
    "    input_feature = FL.data(name='pixel', shape=data_shape, dtype='float32')\n",
    "    hm = FL.data(name='label', shape=data_shape[1:], dtype='float32')\n",
    "    \n",
    "    logit = model.net(input_feature, class_dim=1)\n",
    "    pred_hm = FL.squeeze(\n",
    "        FL.conv2d_transpose(logit, num_filters=1, output_size=256), axes=[1]) # Bs, 256,256\n",
    "    \n",
    "    reader = fluid.io.PyReader(feed_list=[input_feature, hm], \n",
    "                         capacity=64, iterable=True, use_double_buffer=True)\n",
    "\n",
    "    cost = FL.square_error_cost(pred_hm, hm)\n",
    "    loss = FL.mean(cost)\n",
    "    \n",
    "    return [loss, pred_hm, reader]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_dist(pred_hm, hm):\n",
    "    hm = np.array(hm)\n",
    "    \n",
    "    mean_dis = 0.\n",
    "    for single_hm, single_pred_hm in zip(hm, pred_hm):\n",
    "        # Find argmax_x, argmax_y from 2D tensor\n",
    "        label_x, label_y = np.unravel_index(single_hm.argmax(), single_hm.shape)\n",
    "        pred_x, pred_y = np.unravel_index(single_pred_hm.argmax(), single_pred_hm.shape)\n",
    "        mean_dis += np.sqrt((pred_x - label_x) ** 2 + (pred_y - label_y) ** 2)\n",
    "    \n",
    "    return mean_dis / hm.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(use_cuda, params_dirname_prefix, pretrained_model=False, EPOCH_NUM=10):\n",
    "    place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()\n",
    "    \n",
    "    startup_prog = fluid.Program()\n",
    "    train_prog = fluid.Program()\n",
    "    val_prog = fluid.Program()\n",
    "\n",
    "    with fluid.program_guard(train_prog, startup_prog):\n",
    "        # fluid.unique_name.guard() to share parameters with test network\n",
    "        with fluid.unique_name.guard():\n",
    "            train_loss, train_output, train_reader = network()\n",
    "            \n",
    "            optimizer = fluid.optimizer.Adam(learning_rate=1e-4)\n",
    "            optimizer.minimize(train_loss)\n",
    "    \n",
    "    # 定义预测网络\n",
    "    with fluid.program_guard(val_prog, startup_prog):\n",
    "        # Use fluid.unique_name.guard() to share parameters with train network\n",
    "        with fluid.unique_name.guard():\n",
    "            val_loss, val_output, val_reader = network()\n",
    "\n",
    "    val_prog = val_prog.clone(for_test=True)\n",
    "\n",
    "    train_loss.persistable = True\n",
    "    val_loss.persistable = True\n",
    "    val_output.persistable = True\n",
    "            \n",
    "    exe = fluid.Executor(place)\n",
    "    exe.run(startup_prog)\n",
    "\n",
    "    if pretrained_model:\n",
    "        def if_exist(var):\n",
    "            return os.path.exists(os.path.join(pretrained_model, var.name))\n",
    "\n",
    "        fluid.io.load_vars(\n",
    "            exe, pretrained_model, main_program=train_prog, predicate=if_exist)\n",
    "\n",
    "    train_reader.decorate_sample_list_generator( train_dataloader, places=place )\n",
    "    val_reader.decorate_sample_list_generator( val_dataloader, places=place )\n",
    "\n",
    "    # For training test cost\n",
    "    def train_test(val_prog, val_reader):\n",
    "        count = 0\n",
    "        accumulated = [0,0]\n",
    "        \n",
    "        prediction = []\n",
    "        label_values = []\n",
    "        \n",
    "        for tid, val_data in enumerate(val_reader()):\n",
    "            avg_cost_np = exe.run(\n",
    "                program=val_prog,\n",
    "                feed=val_data,\n",
    "                fetch_list=[val_loss, val_output],\n",
    "                use_program_cache=True)\n",
    "            accumulated = [\n",
    "                x[0] + x[1][0] for x in zip(accumulated, avg_cost_np)\n",
    "            ]\n",
    "            prediction.append(avg_cost_np[1])\n",
    "            label_values.append( np.array(val_data[0]['label']) )\n",
    "            count += 1\n",
    "\n",
    "        prediction = np.concatenate(prediction, 0)\n",
    "        label_values = np.concatenate(label_values, 0)\n",
    "\n",
    "        mean_dis = calc_dist(prediction, label_values)\n",
    "        \n",
    "        return [x / count for x in accumulated], mean_dis\n",
    "\n",
    "    # main train loop.\n",
    "    def train_loop():\n",
    "        step = 0\n",
    "        best_dist = 65536.\n",
    "\n",
    "        for pass_id in range(EPOCH_NUM):\n",
    "            data_load_time = time.time()\n",
    "            for step_id, data_train in enumerate(train_reader()):\n",
    "                data_load_costtime = time.time() - data_load_time\n",
    "                start_time = time.time()\n",
    "                avg_loss_value = exe.run(\n",
    "                    train_prog,\n",
    "                    feed=data_train,\n",
    "                    fetch_list=[train_loss, train_output], \n",
    "                    use_program_cache=True)\n",
    "                cost_time = time.time() - start_time\n",
    "                if step_id % 50 == 0:\n",
    "                    mean_dis = calc_dist(avg_loss_value[1], data_train[0]['label'])\n",
    "                    print(\"Pass %d, Epoch %d, Cost %f, EuDis %f, Time %f, LoadTime %f\" % (\n",
    "                        step_id, pass_id, avg_loss_value[0], mean_dis, cost_time, data_load_costtime))\n",
    "                else:\n",
    "                    pass\n",
    "                step += 1\n",
    "                data_load_time = time.time()\n",
    "\n",
    "            avg_cost_test, avg_dist_test = train_test(val_prog, val_reader)\n",
    "\n",
    "            print('Test with Epoch {0}, Loss {1:2.4}, EuDis {2:2.4}'.format(\n",
    "                pass_id, avg_cost_test[0], avg_dist_test))\n",
    "\n",
    "            if avg_dist_test < best_dist:\n",
    "                best_dist = avg_dist_test\n",
    "                print(\"\\nBest Dis, Checkpoint Saved!\\n\")\n",
    "                if not os.path.isdir(params_dirname_prefix+\"_best/\"):\n",
    "                    os.makedirs(params_dirname_prefix+\"_best/\")\n",
    "                fluid.io.save_persistables(exe, params_dirname_prefix+\"_best/\", main_program=train_prog)\n",
    "\n",
    "            if not os.path.isdir(params_dirname_prefix+\"_checkpoint/\"):\n",
    "                os.makedirs(params_dirname_prefix+\"_checkpoint/\")\n",
    "            fluid.io.save_persistables(exe, params_dirname_prefix+\"_checkpoint/\", main_program=train_prog)\n",
    "    train_loop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download imagenet pretrain weight from:\n",
    "# https://github.com/PaddlePaddle/models/tree/develop/PaddleCV/image_classification\n",
    "# remove ResNet34_pretrained/fc*\n",
    "train(use_cuda=True, params_dirname_prefix=\"../weights/loc_fcn\", \n",
    "        pretrained_model=\"../weights/ResNet34_pretrained\", EPOCH_NUM=40)"
   ]
  }
 ],
 "metadata": {
  "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}