post_training_data_cleanup.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use this notebook to load a trained model and use the model to do post procesing data cleanup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import json\n",
    "import h5py\n",
    "import threading\n",
    "from collections import Counter\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "from skimage.io import imread\n",
    "from sklearn.model_selection import train_test_split\n",
    "from skimage.segmentation import find_boundaries\n",
    "from skimage.morphology import binary_opening, disk\n",
    "from skimage.transform import resize\n",
    "from scipy import ndimage as ndi\n",
    "\n",
    "from IPython.display import SVG\n",
    "import tensorflow as tf\n",
    "import keras.backend as K\n",
    "from keras.utils.vis_utils import plot_model\n",
    "from keras.utils.vis_utils import model_to_dot\n",
    "from keras.utils import Sequence\n",
    "from keras.models import load_model\n",
    "from keras.models import Model\n",
    "from keras.layers import Input\n",
    "from keras.layers import Conv2D\n",
    "from keras.layers import Cropping2D\n",
    "from keras.layers import MaxPooling2D\n",
    "from keras.layers import UpSampling2D\n",
    "from keras.layers import ZeroPadding2D\n",
    "from keras.layers import GaussianNoise\n",
    "from keras.layers import Conv2DTranspose\n",
    "from keras.layers import BatchNormalization\n",
    "from keras.layers import Activation\n",
    "from keras.layers import AvgPool2D\n",
    "from keras.layers import Dropout\n",
    "from keras.layers import Flatten\n",
    "from keras.layers import concatenate\n",
    "\n",
    "from keras.optimizers import Adam\n",
    "from keras.losses import binary_crossentropy\n",
    "from keras.preprocessing.image import ImageDataGenerator\n",
    "from keras.preprocessing.image import Iterator\n",
    "\n",
    "from utils import get_datetime_now\n",
    "from utils import reversed_recombined_holt_winters\n",
    "from utils import rle_encode\n",
    "from utils import rle_decode\n",
    "from utils import multi_rle_encode\n",
    "from utils import masks_as_image\n",
    "from utils import get_mask_weights\n",
    "\n",
    "_EPSILON = K.epsilon()\n",
    "\n",
    "SMALL_SIZE = 10\n",
    "MEDIUM_SIZE = 12\n",
    "BIGGER_SIZE = 16\n",
    "BIGGEST_SIZE = 20\n",
    "plt.rc('font', size=BIGGEST_SIZE)         # controls default text sizes\n",
    "plt.rc('axes', titlesize=BIGGEST_SIZE)    # fontsize of the axes title\n",
    "plt.rc('axes', labelsize=BIGGEST_SIZE)    # fontsize of the x and y labels\n",
    "plt.rc('xtick', labelsize=BIGGER_SIZE)   # fontsize of the tick labels\n",
    "plt.rc('ytick', labelsize=BIGGER_SIZE)   # fontsize of the tick labels\n",
    "plt.rc('legend', fontsize=MEDIUM_SIZE)   # legend fontsize\n",
    "plt.rc('figure', titlesize=BIGGEST_SIZE)  # fontsize of the figure title"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dir_ext = ''\n",
    "BASE_MODEL = \"2018_1107_1009_43_GCN_768_focaldice_deconv\"\n",
    "NORMALIZE_INPUT_BATCH = True\n",
    "GAUSSIAN_NOISE = None\n",
    "DROPOUT = 0\n",
    "LEARN_RATE = 1e-3\n",
    "# LOSS_FUNCTION = \"bce_soft_dice_loss\"\n",
    "LOSS_FUNCTION = \"focal_soft_dice_loss\"\n",
    "LOSS_FUNCTION_KWARGS = {\"gamma\": 2.0, \"alpha\": 0.25}\n",
    "MASK_METHOD = \"fullmasks\"\n",
    "IMG_SIZE = 768  # 768, 384, 192, 96\n",
    "CROP_SIZE = 768  # 512, 256, 128, 64\n",
    "EDGE_CROP = (IMG_SIZE - CROP_SIZE) // 8\n",
    "BATCH_SIZE = 8\n",
    "# INPUT_SHAPE = (CROP_SIZE, CROP_SIZE, 3)\n",
    "INPUT_SHAPE = (None, None, 3)\n",
    "# INPUT_SHAPE = (221, 221, 3)\n",
    "# number of training images to use (-1 == use all)\n",
    "TRAIN_IMG_COUNT = -1\n",
    "# number of validation images to use (-1 == use all)\n",
    "VALID_IMG_COUNT = -1\n",
    "# Dict for saving all configuration data.\n",
    "data_config = {\n",
    "    'base_model': BASE_MODEL,\n",
    "    'normalize_input_batch': NORMALIZE_INPUT_BATCH,\n",
    "    'gaussian_noise': GAUSSIAN_NOISE,\n",
    "    'dropout': DROPOUT,\n",
    "    'learn_rate': LEARN_RATE,\n",
    "    'loss_function': LOSS_FUNCTION,\n",
    "    'loss_function_kwargs': LOSS_FUNCTION_KWARGS,\n",
    "    'mask_method': MASK_METHOD,\n",
    "    'img_size': IMG_SIZE,\n",
    "    'crop_size': CROP_SIZE,\n",
    "    'edge_crop': EDGE_CROP,\n",
    "    'batch_size': BATCH_SIZE,\n",
    "    'input_shape': INPUT_SHAPE,\n",
    "    'train_img_count': TRAIN_IMG_COUNT,\n",
    "    'valid_img_count': VALID_IMG_COUNT,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from skimage.util import montage\n",
    "montage_rgb = lambda x: np.stack([montage(x[:, :, :, i]) for i in range(x.shape[3])], -1)\n",
    "# montage_rgb = lambda x: np.stack([montage(x[:, :, :, i], padding_width=10, fill=[255, 255, 255]) for i in range(x.shape[3])], -1)\n",
    "montage_pad = lambda x, *args, **kwargs: montage(x, padding_width=10, *args, **kwargs)\n",
    "ship_dir = \"/media/Borg_LS/DATA/geos/airbus/input/\"\n",
    "train_img_path = os.path.join(ship_dir, f'train_{IMG_SIZE}')\n",
    "train_seg_file = os.path.join(ship_dir, f\"{MASK_METHOD}_{IMG_SIZE}.h5\")\n",
    "train_meta_filebase = f\"{MASK_METHOD}_{IMG_SIZE}\"\n",
    "# train_meta_filebase = f\"{MASK_METHOD}_768_down_{IMG_SIZE}\"\n",
    "# train_meta_filebase = \"background_stats\"\n",
    "train_meta_file = os.path.join(ship_dir, train_meta_filebase + '.csv')\n",
    "datetime_now = get_datetime_now()\n",
    "run_name = f\"{datetime_now}_{dir_ext}\"\n",
    "out_dir = f\"out/{run_name}/\"\n",
    "\n",
    "# image_weights = pd.read_csv(background_stats_file)\n",
    "train_df = pd.read_csv(train_meta_file)\n",
    "print(train_meta_filebase + '.csv')\n",
    "print(train_df.shape[0], 'records found')\n",
    "print(train_df['ImageId'].value_counts().shape[0])\n",
    "train_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# We make a generator to produce batches of images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     151,
     158,
     163,
     170,
     212,
     218,
     228
    ]
   },
   "outputs": [],
   "source": [
    "# from utils import fbeta\n",
    "\n",
    "class DataGenerator(Sequence):\n",
    "    \"\"\"Generates data for Keras\"\"\"\n",
    "    def __init__(self,\n",
    "                 img_df,\n",
    "                 img_dir,\n",
    "                 seg_file,\n",
    "                 img_size=768,\n",
    "                 crop_size=None,\n",
    "                 edge_crop=None,\n",
    "                 batch_size=4,\n",
    "                 n_channels=3,\n",
    "                 n_classes=1,\n",
    "                 n_samples=-1,\n",
    "                 use_weights=True,\n",
    "                 shuffle=True,\n",
    "                 ig_args=None,\n",
    "                 dg_args=None,\n",
    "                 seed=None\n",
    "                 ):\n",
    "\n",
    "        \"\"\"Initialization\"\"\"\n",
    "#         from utils import fbeta\n",
    "#         from utils import focal_soft_dice_loss_wrapper\n",
    "#         self.weight_func = focal_soft_dice_loss_wrapper(focal_coef=0.2, axis=None)\n",
    "#         from utils import soft_dice_loss\n",
    "#         self.weight_func = soft_dice_loss\n",
    "#         from utils import focal_loss_wrapper\n",
    "#         self.weight_func = focal_loss_wrapper(axis=None)\n",
    "#         from utils import f2score\n",
    "#         self.weight_func = f2score\n",
    "#         self.weight_func = mean_fscore\n",
    "#         self.weight_func = lambda x, y: np.sum((x - y)**2, axis=(1, 2, 3)) / np.prod(x.shape[1:])  # briar loss\n",
    "        self.weight_func = lambda x, y: np.sum((x - y)**2) / np.prod(x.shape)  # briar loss\n",
    "        self.lock = threading.Lock()\n",
    "        \n",
    "        if 'ImageId0' in img_df:\n",
    "            self.img0_ids = img_df['ImageId0'].values\n",
    "            \n",
    "        self.img_ids = img_df['ImageId'].values\n",
    "        self.n_records = len(self.img_ids)\n",
    "        if isinstance(img_size, tuple):\n",
    "            assert img_size[0] == img_size[1]\n",
    "            self.img_size = img_size[0]\n",
    "        else:\n",
    "            self.img_size = img_size\n",
    "            \n",
    "        if crop_size is None:\n",
    "            crop_size = img_size\n",
    "        if isinstance(crop_size, tuple):\n",
    "            assert crop_size[0] == crop_size[1]\n",
    "            self.crop_size = crop_size[0]\n",
    "        else:\n",
    "            self.crop_size = crop_size\n",
    "            \n",
    "        assert self.crop_size <= self.img_size\n",
    "        \n",
    "        if edge_crop is None:\n",
    "            self.edge_crop = (self.img_size - self.crop_size) // 8\n",
    "            \n",
    "        self.counts = img_df['counts'].values\n",
    "        self.all_counts = np.sum(img_df['counts'].values)\n",
    "        self.norm_counts = img_df['counts'].values / self.all_counts\n",
    "        \n",
    "#         self.pixels = img_df['n_pixels'].values\n",
    "#         self.all_pixels = np.sum(img_df['n_pixels'].values)\n",
    "#         self.norm_pixels = img_df['n_pixels'].values / self.all_pixels\n",
    "        \n",
    "        if 'weight' in img_df:\n",
    "            self.img_weights = img_df['weight'].values\n",
    "            self.all_weight = np.sum(self.img_weights)\n",
    "            self.norm_weights = self.img_weights / self.all_weight\n",
    "        else: \n",
    "            self.use_weights = False\n",
    "\n",
    "        self.norm_weights = np.ones(self.n_records)\n",
    "        self.loss_weights = np.ones(self.n_records)\n",
    "        self.use_weights = use_weights\n",
    "        \n",
    "        self.img_dir = img_dir\n",
    "        self.seg_file = seg_file\n",
    "        self.shuffle = shuffle\n",
    "        self.n_channels = n_channels\n",
    "        self.n_classes = n_classes\n",
    "\n",
    "        self._n_samples = self.n_records if n_samples in (None, -1) else min(n_samples, self.n_records)\n",
    "        self._batch_size = self._n_samples if batch_size in (None, -1) else min(batch_size, self._n_samples)\n",
    "\n",
    "        self.access_counter = Counter()\n",
    "#         self.epochs_seen = 0\n",
    "#         self.images_seen = 0\n",
    "#         self.weight_gain = 0\n",
    "#         self.ships_seen = 0\n",
    "\n",
    "        self.augment = True\n",
    "        if dg_args is not None and ig_args is None:\n",
    "            self.image_gen = ImageDataGenerator(**dg_args)\n",
    "            self.label_gen = ImageDataGenerator(**dg_args)\n",
    "        elif dg_args is None and ig_args is not None: # Validation\n",
    "            self.image_gen = ImageDataGenerator(**ig_args)\n",
    "            self.label_gen = ImageDataGenerator()\n",
    "        elif dg_args is not None and ig_args is not None:\n",
    "            ig_args = dict(**ig_args, **dg_args)\n",
    "            self.image_gen = ImageDataGenerator(**ig_args)\n",
    "            self.label_gen = ImageDataGenerator(**dg_args)\n",
    "        else:\n",
    "            self.augment = False  # to get x_train_mean\n",
    "        \n",
    "        np.random.seed(seed if seed is not None else np.random.choice(range(9999)))\n",
    "        self.on_epoch_end()\n",
    "\n",
    "#     @property\n",
    "#     def ships_score(self):\n",
    "#         return self.ships_seen / self.epochs_seen / self.sample_fraction\n",
    "    \n",
    "#     @property\n",
    "#     def weight_score(self):\n",
    "#         return self.weight_gain / self.epochs_seen / self.sample_fraction\n",
    "    \n",
    "    @property\n",
    "    def n_samples(self):\n",
    "        return self._n_samples\n",
    "            \n",
    "    @n_samples.setter\n",
    "    def n_samples(self, n_samples):\n",
    "        self._n_samples = self.n_records if n_samples in (None, -1) else min(n_samples, self.n_records)\n",
    "        self._batch_size = min(self._batch_size, self._n_samples)\n",
    "        \n",
    "    @property\n",
    "    def sample_fraction(self):\n",
    "        self.n_samples / self.n_records\n",
    "        \n",
    "    @property\n",
    "    def batch_size(self):\n",
    "        return self._batch_size\n",
    "    \n",
    "    @batch_size.setter\n",
    "    def batch_size(self, batch_size):\n",
    "        self._batch_size = self.n_samples if batch_size in (None, -1) else min(batch_size, self.n_samples)\n",
    "        \n",
    "    def update_record_weights(self, y_true, y_hat):\n",
    "#         y_pred = (y_hat >= 0.5) * 1\n",
    "        \n",
    "#         sample_losses = self.weight_func(y_true, y_hat)\n",
    "#         self.loss_weights[self.current_index_array] = sample_losses\n",
    "        \n",
    "        for i, idx in enumerate(self.current_index_array):\n",
    "            # The score for each individual sample in current_index_array.\n",
    "            \n",
    "            # subtract weight_func from 1 if weight_func is a metric (e.g. acc, IOU, Fbeta, etc.)\n",
    "#             sample_scores = self.weight_func(y_true[i], y_pred[i])\n",
    "#             sample_losses = 1 - sample_scores\n",
    "\n",
    "            sample_losses = self.weight_func(y_true[i], y_hat[i])\n",
    "    \n",
    "            self.loss_weights[idx] = sample_losses\n",
    "#             self.access_counter[idx] += 1\n",
    "\n",
    "        # Setting index_array to None reminds us that we must call update_record_weights before on_epoch_end.\n",
    "#         self.index_array = None\n",
    "#         self.on_epoch_end()\n",
    "\n",
    "    def print_weight_summary(self):\n",
    "        for idx in range(self.n_records):\n",
    "            if idx in self.index_array:\n",
    "                print(f\"{idx:>3} {self.access_counter[idx]:>3} {self.loss_weights[idx]:>.3f} {self.norm_weights[idx]:>.3f} index_array\")\n",
    "            else:\n",
    "                print(f\"{idx:>3} {self.access_counter[idx]:>3} {self.loss_weights[idx]:>.3f} {self.norm_weights[idx]:>.3f}\")\n",
    "            \n",
    "    def __len__(self):\n",
    "        \"\"\"Denotes the number of batches per epoch\"\"\"\n",
    "        return (self.n_samples + self.batch_size - 1) // self.batch_size  # round up\n",
    "#         return self.n_samples // self.batch_size\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        \"\"\"Generate one batch of data\"\"\"\n",
    "        # Generate indexes of the batch\n",
    "        self.current_index_array = self.index_array[self.batch_size * idx: self.batch_size * (idx + 1)]\n",
    "        # Generate and return batch\n",
    "        return self._get_batches_of_transformed_samples(self.current_index_array)\n",
    "\n",
    "    def next(self):\n",
    "        \"\"\"For python 2.x.\n",
    "\n",
    "        # Returns\n",
    "            The next batch.\n",
    "        \"\"\"\n",
    "        # Keeps under lock only the mechanism which advances\n",
    "        # the indexing of each batch.\n",
    "        with self.lock:\n",
    "            index_array = next(self.index_generator)\n",
    "        # The transformation of images is not under thread lock\n",
    "        # so it can be done in parallel\n",
    "        return self._get_batches_of_transformed_samples(index_array)\n",
    "\n",
    "    def on_epoch_end(self):\n",
    "        \"\"\"Updates indexes after each epoch\"\"\"\n",
    "        if self.n_samples < self.n_records:\n",
    "            if self.use_weights:\n",
    "                self.norm_weights = self.loss_weights / np.sum(self.loss_weights)\n",
    "                self.index_array = np.random.choice(self.n_records, self.n_samples, replace=False, p=self.norm_weights)\n",
    "            else:\n",
    "                self.index_array = np.random.choice(self.n_records, self.n_samples, replace=False, p=None)\n",
    "        else:\n",
    "            self.index_array = np.arange(self.n_records)\n",
    "#         self.epochs_seen += 1\n",
    "        if self.shuffle:\n",
    "            np.random.shuffle(self.index_array)\n",
    "\n",
    "    def get_next_batch(self, index_array):\n",
    "        \"\"\"Generates data containing batch_size samples\"\"\"\n",
    "        x_batch = np.empty((len(index_array), self.img_size, self.img_size, self.n_channels), dtype=np.float32)\n",
    "        y_batch = np.empty((len(index_array), self.img_size, self.img_size, self.n_classes), dtype=np.float32)\n",
    "        with h5py.File(self.seg_file, 'r') as mask_1:\n",
    "            for i, idx in enumerate(index_array):\n",
    "                x_batch[i] = cv2.cvtColor(cv2.imread(os.path.join(self.img_dir, self.img_ids[idx])), cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0\n",
    "                try:\n",
    "                    y_batch[i] = mask_1[self.img_ids[idx]][:] * 1\n",
    "                except KeyError:\n",
    "                    y_batch[i] = np.zeros((self.img_size, self.img_size, self.n_classes), dtype=np.float32)\n",
    "        return x_batch, y_batch\n",
    "    \n",
    "    def augment_batch(self, x_batch, y_batch):\n",
    "        seed = np.random.choice(range(9999))\n",
    "        x_batch_aug = self.image_gen.flow(x_batch, batch_size=self.batch_size, seed=seed, shuffle=False).next()\n",
    "        y_batch_aug = self.label_gen.flow(y_batch, batch_size=self.batch_size, seed=seed, shuffle=False).next()\n",
    "        return x_batch_aug, (y_batch_aug > 0.5) * 1.0 # make sure mask is still zeros and ones.\n",
    "    \n",
    "    def crop_batch(self, x_batch_aug, y_batch_aug):\n",
    "        x_batch_aug_crops = np.empty((self.batch_size, self.crop_size, self.crop_size, self.n_channels), dtype=np.float32)\n",
    "        y_batch_aug_crops = np.empty((self.batch_size, self.crop_size, self.crop_size, self.n_classes), dtype=np.float32)\n",
    "        for i in range(self.batch_size):\n",
    "            x = np.random.randint(self.edge_crop, self.img_size - self.crop_size - self.edge_crop + 1)\n",
    "            y = np.random.randint(self.edge_crop, self.img_size - self.crop_size - self.edge_crop + 1)\n",
    "            x_batch_aug_crops[i] = x_batch_aug[i, y:(y + self.crop_size), x:(x + self.crop_size), :]\n",
    "            y_batch_aug_crops[i] = y_batch_aug[i, y:(y + self.crop_size), x:(x + self.crop_size), :]\n",
    "        return x_batch_aug_crops, y_batch_aug_crops\n",
    "        \n",
    "    def _get_batches_of_transformed_samples(self, index_array):\n",
    "        \n",
    "        # Create the raw batches\n",
    "        x_batch, y_batch = self.get_next_batch(index_array)\n",
    "        # Augment\n",
    "        if self.augment:\n",
    "            x_batch, y_batch = self.augment_batch(x_batch, y_batch)\n",
    "        # Crop\n",
    "        if self.crop_size < self.img_size:\n",
    "            x_batch, y_batch = self.crop_batch(x_batch, y_batch)\n",
    "        \n",
    "        return x_batch, y_batch\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def max_rotation(edge_crop=EDGE_CROP, img_size=IMG_SIZE, crop_size=CROP_SIZE):\n",
    "    return np.arccos(np.sqrt((crop_size + 2 * edge_crop) / img_size)) * 180.0 / np.pi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# These args are applied to both training inputs and training outputs.\n",
    "dg_args = dict(\n",
    "#     height_shift_range = 32, # If random cropping, don't need to shift\n",
    "#     width_shift_range = 32, # If random cropping, don't need to shift\n",
    "#     shear_range = 0.01,\n",
    "#     zoom_range = [0.9, 1.25], \n",
    "    horizontal_flip = True, \n",
    "    vertical_flip = True,\n",
    "#     rotation_range = max_rotation(), \n",
    "    fill_mode = 'reflect',\n",
    "    data_format = 'channels_last'\n",
    ")\n",
    "\n",
    "train_gen = DataGenerator(\n",
    "    train_df, \n",
    "    train_img_path, \n",
    "    train_seg_file, \n",
    "    img_size=IMG_SIZE, \n",
    "    crop_size=CROP_SIZE,\n",
    "    batch_size=TRAIN_IMG_COUNT, \n",
    "    n_samples=1024, \n",
    "    use_weights=True,\n",
    "    shuffle=False,\n",
    "    seed=777, \n",
    "#     dg_args=dg_args, \n",
    "    ig_args=None)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "brier_loss = lambda x, y: np.mean(np.sum((x - y)**2, axis=-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fbeta(y_true, y_pred, beta=2.0, axis=-1, smooth=1e-3):\n",
    "    tp = K.sum(y_true * y_pred, axis=axis)\n",
    "    fp = K.sum((1 - y_true) * y_pred, axis=axis)\n",
    "    fn = K.sum(y_true * (1 - y_pred), axis=axis)\n",
    "    return K.mean((beta**2 + 1) * tp / ((beta**2 + 1) * tp + beta**2 * fn + fp + smooth))\n",
    "\n",
    "# CANT USE K.round IN LOSS FUNCTION. NOT DIFFERENTIABLE\n",
    "def hard_dice_coef(y_true, y_pred, axis=-1, smooth=1e-3):\n",
    "    return soft_dice_coef(K.round(y_true), K.round(y_pred), axis=axis, smooth=smooth)\n",
    "\n",
    "def soft_dice_coef(y_true, y_pred, axis=-1, smooth=1e-3):\n",
    "    AB = K.sum(y_true * y_pred, axis=axis)\n",
    "    A = K.sum(y_true, axis=axis)\n",
    "    B = K.sum(y_pred, axis=axis)\n",
    "    return (2. * AB + smooth) / (A + B + smooth)\n",
    "\n",
    "def soft_dice_loss(y_true, y_pred, axis=-1, smooth=1e-3):\n",
    "    return 1 - soft_dice_coef(y_true, y_pred, axis=axis, smooth=smooth)\n",
    "\n",
    "def bce_soft_dice_loss_wrapper(bce_weight=0.5):\n",
    "    def bce_soft_dice_loss(y_true, y_pred):\n",
    "        return binary_crossentropy(y_true, y_pred) * bce_weight + soft_dice_loss(y_true, y_pred, axis=-1) * (1 - bce_weight)\n",
    "    return bce_soft_dice_loss\n",
    "\n",
    "def focal_loss_wrapper(gamma=2., alpha=.25, axis=-1):\n",
    "    def focal_loss(y_true, y_pred):\n",
    "        y_pred_c = K.clip(y_pred, _EPSILON, 1.0 - _EPSILON)\n",
    "        pt_1 = tf.where(tf.equal(y_true, 1), y_pred_c, tf.ones_like(y_pred))\n",
    "        pt_0 = tf.where(tf.equal(y_true, 0), y_pred_c, tf.zeros_like(y_pred))\n",
    "        res1 =      alpha  * K.pow(1. - pt_1, gamma) * K.log(     pt_1)\n",
    "        res0 = (1 - alpha) * K.pow(     pt_0, gamma) * K.log(1. - pt_0)\n",
    "        return -K.mean(res1 + res0, axis=axis)\n",
    "    return focal_loss\n",
    "\n",
    "def focal_soft_dice_loss_wrapper(gamma=2., alpha=.25, axis=-1):\n",
    "    focal_loss = focal_loss_wrapper(gamma=gamma, alpha=alpha, axis=axis)\n",
    "    def focal_soft_dice_loss(y_true, y_pred, axis=-1):\n",
    "        return focal_loss(y_true, y_pred) * 0.8 + soft_dice_loss(y_true, y_pred, axis=axis) * 0.2\n",
    "    return focal_soft_dice_loss\n",
    "\n",
    "def brier_loss_keras(y_true, y_pred):\n",
    "    return K.sum((y_pred - y_true)**2, axis=-1)\n",
    "\n",
    "loss_functions = {\n",
    "    \"focal_loss\": focal_loss_wrapper,\n",
    "    \"soft_dice_loss\": soft_dice_loss,\n",
    "    \"bce_soft_dice_loss\": bce_soft_dice_loss_wrapper,\n",
    "    \"focal_soft_dice_loss\": focal_soft_dice_loss_wrapper\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import get_best_model_name\n",
    "best_model_name = get_best_model_name(BASE_MODEL)\n",
    "seg_model = load_model(best_model_name, compile=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "focal_loss = focal_loss_wrapper()\n",
    "seg_model.compile(\n",
    "    optimizer=Adam(lr=LEARN_RATE), \n",
    "#     loss=brier_loss_keras,\n",
    "    loss=loss_functions[LOSS_FUNCTION](**LOSS_FUNCTION_KWARGS),\n",
    "    metrics=[binary_crossentropy, focal_loss, brier_loss_keras])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "n_samples_temp = train_gen.n_samples\n",
    "print(len(train_gen), train_gen.n_samples, len(train_gen.index_array), train_gen.batch_size, n_samples_temp)\n",
    "train_gen.n_samples = -1\n",
    "print(len(train_gen), train_gen.n_samples, len(train_gen.index_array), train_gen.batch_size, n_samples_temp)\n",
    "train_gen.on_epoch_end()\n",
    "print(len(train_gen), train_gen.n_samples, len(train_gen.index_array), train_gen.batch_size, n_samples_temp)\n",
    "\n",
    "for i in range(len(train_gen)):\n",
    "    t0 = time.time()\n",
    "    x_train, y_train = train_gen[i]\n",
    "    t1 = time.time()\n",
    "    y_hat = seg_model.predict(x_train, batch_size=BATCH_SIZE, verbose=1)\n",
    "    t2 = time.time()\n",
    "    train_gen.update_record_weights(y_train, y_hat)\n",
    "    print(\n",
    "        f\"{t1 - t0:>.3f} {t2 - t1:>.3f} {time.time() - t2:>.3f}\", \n",
    "        f\"{np.min(train_gen.loss_weights):>9.5e} {np.max(train_gen.loss_weights):>9.5g}\", \n",
    "        f\"{np.mean(train_gen.loss_weights):>9.5f} {np.median(train_gen.loss_weights):>9.5g} {np.std(train_gen.loss_weights):>9.5f}\"\n",
    "    )\n",
    "    \n",
    "print(len(train_gen), train_gen.n_samples, len(train_gen.index_array), train_gen.batch_size, n_samples_temp)\n",
    "train_gen.n_samples = n_samples_temp\n",
    "print(len(train_gen), train_gen.n_samples, len(train_gen.index_array), train_gen.batch_size, n_samples_temp)\n",
    "train_gen.on_epoch_end()\n",
    "print(len(train_gen), train_gen.n_samples, len(train_gen.index_array), train_gen.batch_size, n_samples_temp)\n",
    "loss_weight_indices = np.argsort(train_gen.loss_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for i, layer in enumerate(seg_model.layers):\n",
    "    print(i+1, layer.name, layer.trainable)\n",
    "    weights = layer.get_weights()\n",
    "    for w in weights:\n",
    "        print('  ', w.shape, w.min(), w.max())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loss_weights_old = train_df['brier'].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(np.sum(train_loss_weights_old), np.sum(train_gen.loss_weights), np.sum(train_gen.loss_weights - train_loss_weights_old))\n",
    "len(np.where(train_loss_weights_old > train_gen.loss_weights)[0]) / len(train_loss_weights_old)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_samples = 8\n",
    "new_worst_idx = loss_weight_indices[-n_samples:][::-1]\n",
    "new_worst_gen_loss = train_gen.loss_weights[new_worst_idx]\n",
    "new_worst_gen_ids = train_gen.img_ids[new_worst_idx]\n",
    "new_worst_df_loss = train_df['brier'].values[new_worst_idx]\n",
    "new_worst_df_ids = train_df['ImageId'].values[new_worst_idx]\n",
    "new_header = (' ' * len(new_worst_gen_ids[0]), 'new worst', 'prev loss', 'change')\n",
    "print(f\"{new_header[0]} {new_header[1]:>10} {new_header[2]:>10} {new_header[3]:>10}\")\n",
    "for wgl, wgi, wdl, wdi in zip(new_worst_gen_loss, new_worst_gen_ids, new_worst_df_loss, new_worst_df_ids):\n",
    "    assert wgi == wdi\n",
    "    print(f'{wgi} {wgl:<10.6g} {wdl:<10.6g} {wgl - wdl:<10.6g}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "old_loss_weight_indices = np.argsort(train_df['brier'].values)\n",
    "old_worst_idx = old_loss_weight_indices[-8:][::-1]\n",
    "old_worst_df_loss = train_df['brier'].values[old_worst_idx]\n",
    "old_worst_df_ids = train_df['ImageId'].values[old_worst_idx]\n",
    "old_worst_gen_loss = train_gen.loss_weights[old_worst_idx]\n",
    "old_worst_gen_ids = train_gen.img_ids[old_worst_idx]\n",
    "old_header = (' ' * len(old_worst_df_ids[0]), 'prev worst', 'new loss', 'change')\n",
    "print(f\"{old_header[0]} {old_header[1]:>10} {old_header[2]:>10} {old_header[3]:>10}\")\n",
    "for wdl, wdi, wgl, wgi in zip(old_worst_df_loss, old_worst_df_ids, old_worst_gen_loss, old_worst_gen_ids):\n",
    "    assert wgi == wdi\n",
    "    print(f'{wdi} {wdl:<10.6g} {wgl:<10.6g} {wgl - wdl:<10.6g}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fig, m_axs = plt.subplots(n_samples, 3, figsize = (16, 42))\n",
    "with h5py.File(train_seg_file, 'r') as mask_1:\n",
    "    for (ax1, ax2, ax4), c_img_id, c_img_loss in zip(m_axs, old_worst_df_ids, old_worst_df_loss):\n",
    "        c_file = os.path.join(train_img_path, c_img_id)\n",
    "        c_img = cv2.cvtColor(cv2.imread(c_file), cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0\n",
    "        c_seg = seg_model.predict(c_img[None, ...])\n",
    "        cur_seg = binary_opening(c_seg[0] > 0.5, np.expand_dims(disk(2), -1))\n",
    "        cur_rles = multi_rle_encode(cur_seg)\n",
    "        try:\n",
    "            y_true = mask_1[c_img_id][:] * 1\n",
    "            print('1', y_true.dtype, y_true.shape, brier_loss(y_true, c_seg))\n",
    "        except KeyError:\n",
    "            y_true = np.zeros((IMG_SIZE, IMG_SIZE, 1), dtype=np.uint8)\n",
    "            print('2', y_true.dtype, y_true.shape, brier_loss(y_true, c_seg))\n",
    "        score = seg_model.evaluate(c_img[None, ...], y_true[None, ...])\n",
    "        \n",
    "        ax1.imshow(c_img)\n",
    "        ax1.set_title(c_img_id)\n",
    "        ax2.imshow(c_seg[0, :, :, 0], vmin = 0, vmax = 1)\n",
    "        ax2.set_title('  '.join(map(lambda x: f\"{x:.3f}\", score)))\n",
    "        ax4.imshow(c_img)\n",
    "        ax4.imshow(y_true[..., -1], vmin = 0, vmax = 1, alpha=0.3)\n",
    "        ax4.set_title(f'brier: {c_img_loss}')\n",
    "\n",
    "plt.tight_layout()\n",
    "fig.savefig(os.path.join('out', BASE_MODEL, f\"{train_meta_filebase}_prev_worst_improved.jpg\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gamma = 0.8\n",
    "train_loss_weights_new = train_loss_weights_old * gamma + train_gen.loss_weights * (1 - gamma)\n",
    "train_loss_weights_old[-3:], train_gen.loss_weights[-3:], train_loss_weights_new[-3:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df['brier'] = train_loss_weights_new"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df['brier'] = train_gen.loss_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df['brier'].hist(bins=100);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = os.path.join(ship_dir, \"background_stats.csv\")\n",
    "train_df = pd.read_csv(fname)\n",
    "print(len(train_df))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_ids_old = train_df['ImageId'].values\n",
    "img_ids_new = np.array([img_id.replace('_768.jpg', '.jpg') for img_id in img_ids_old])\n",
    "train_df['ImageId'] = pd.Series(img_ids_new, index=train_df.index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df = train_df.drop(columns=['ImageId0'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df.rename(columns={'ClusterId_100': 'ClusterId'}, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df.rename(columns={'r_mean': 'b_mean0'}, inplace=True)\n",
    "train_df.rename(columns={'b_mean': 'r_mean'}, inplace=True)\n",
    "train_df.rename(columns={'b_mean0': 'b_mean'}, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df.rename(columns={'r_std': 'b_std0'}, inplace=True)\n",
    "train_df.rename(columns={'b_std': 'r_std'}, inplace=True)\n",
    "train_df.rename(columns={'b_std0': 'b_std'}, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for col in ['r_mean', 'g_mean', 'b_mean']:\n",
    "    val_old = train_df[col].values\n",
    "    val_new = val_old * 255.0\n",
    "    train_df[col] = pd.Series(val_new, index=train_df.index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for col in ['r_std', 'g_std', 'b_std']:\n",
    "    val_old = train_df[col].values\n",
    "    val_new = val_old * 255.0\n",
    "    train_df[col] = pd.Series(val_new, index=train_df.index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# removed from fullmasks_192.csv\n",
    "low_res = [\n",
    "    '1a56132bc_1.jpg', \n",
    "    '1fcd09fdf_1.jpg', \n",
    "    '21ad7a98d_1.jpg', \n",
    "    '221d429af_1.jpg', \n",
    "    '27ded27a6_1.jpg', \n",
    "    '2cd9f6542_1.jpg', \n",
    "    '2d4a3ad9f_1.jpg', \n",
    "    '301d489d9_1.jpg', \n",
    "    '33e019433_1.jpg', \n",
    "    '3d7e3902d_1.jpg', \n",
    "    '3de5abd07_1.jpg', \n",
    "    '4c03a74fa_1.jpg',\n",
    "    '42c90bbc2_1.jpg', \n",
    "    '45365ba8d_1.jpg', \n",
    "    '4eca93831_1.jpg', \n",
    "    '5c1f413c3_1.jpg', \n",
    "    '5cc0ed875_1.jpg', \n",
    "    '66de7e8d6_2.jpg', \n",
    "    '6c85bd549_1.jpg', \n",
    "    '7178a5302_1.jpg', \n",
    "    '7573b6fe0_1.jpg', \n",
    "    '95340daf5_1.jpg', \n",
    "    'a5be35c12_1.jpg', \n",
    "    'abefabc92_1.jpg', \n",
    "    'd19cce902_1.jpg', \n",
    "    'c15b6af5a_1.jpg', \n",
    "    'cb7499731_1.jpg', \n",
    "    'e8f8275fc_1.jpg', \n",
    "    'f0ce68058_1.jpg', \n",
    "    'f1832f81c_1.jpg', \n",
    "    'fed3a4375_1.jpg', \n",
    "]\n",
    "missing_ships = ['a280fe312_1.jpg']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# removed from fullmasks_384_down_192.csv\n",
    "low_res = [\n",
    "    '67c6135dd_1_384.jpg',\n",
    "    '27ded27a6_1_384.jpg',\n",
    "    'f1832f81c_1_384.jpg',\n",
    "    'dce9a2b8c_1_384.jpg']\n",
    "\n",
    "bad_mask = ['cc9271ef5_1_384.jpg']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#bad_overlap\n",
    "low_res = ['4bfc4cc05_1.jpg']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# removed from fullmasks_384.csv\n",
    "low_res = [\n",
    "    '01b030c4b_1.jpg', '1fcd09fdf_1.jpg', \n",
    "    '27ded27a6_1.jpg', '5bf23faa1_1.jpg',\n",
    "    '5c1f413c3_1.jpg', '67c6135dd_1.jpg',\n",
    "    '6c85bd549_1.jpg', '72ee20da7_1.jpg',\n",
    "    '9a67345e4_1.jpg', 'ab6dd064f_1.jpg',\n",
    "    'b7643cfd8_1.jpg', 'b60afd050_1.jpg',\n",
    "    'b9c532adf_1.jpg', 'bf0820f3c_1.jpg',\n",
    "    'c15b6af5a_1.jpg', 'c9fbdab2b_1.jpg',\n",
    "    'cac0ef25b_1.jpg', 'cdea345ac_1.jpg',\n",
    "    'd538b9285_1.jpg', 'dce9a2b8c_1.jpg', \n",
    "    'e8f8275fc_1.jpg', 'e99c164e4_1.jpg',\n",
    "    'f0ce68058_1.jpg', 'f1832f81c_1.jpg',\n",
    "    'fe5b2b445_1.jpg', 'fed3a4375_1.jpg',\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "missing_segs = [\n",
    "    'c98adae25.jpg', 'ae3dc4f96.jpg', 'db6bd5497.jpg',\n",
    "    '09c49fb3d.jpg', 'b98fadbe6.jpg', '261af847a.jpg',\n",
    "    '1798f81c9.jpg', '9984793d9.jpg', '2d80cdcd2.jpg',\n",
    "    '4bef9d21b.jpg', '2410e5ca2.jpg', 'd88bfdafc.jpg', \n",
    "    '738e99b63.jpg', '99ef3e71b.jpg', '731f504e5.jpg',\n",
    "    'd5102732f.jpg', '919e341b8.jpg', '7160d01dc.jpg', \n",
    "    '1b4a6539a.jpg', '99248ac6e.jpg', 'f9e66a6fe.jpg', \n",
    "    'fdf018a2d.jpg', '8a7abc7a2.jpg', 'f331e6546.jpg', \n",
    "    '4d4a1f002.jpg', '3c814ab2e.jpg', 'c9fa3e0b1.jpg', \n",
    "    'f68a4249e.jpg', '4f02b2381.jpg', '9c15e2643.jpg',\n",
    "    '9d2de1ba9.jpg', '88efbddbc.jpg', 'a9364256c.jpg',\n",
    "    '0f30940aa.jpg', 'd6d090131.jpg', '8639577a6.jpg', \n",
    "    'abd0aa2e8.jpg', '266be9837.jpg', '84b35bd43.jpg',\n",
    "    'cebe251ed.jpg', '2da5843a3.jpg', '2cc7799b4.jpg',\n",
    "    'a87026497.jpg', '836e44cc8.jpg', '78376afae.jpg', \n",
    "    'cab1ddb95.jpg', 'ccb365fca.jpg', 'b0c2b7382.jpg',\n",
    "    'ca75a6207.jpg', 'a6bd394af.jpg', '5cbb206f0.jpg', \n",
    "    'c02e5255f.jpg', 'db66e1593.jpg', '390542ac0.jpg', \n",
    "    '3ea817c2b.jpg', 'aab5f1a09.jpg', 'c6b1facda.jpg',\n",
    "    '458bac358.jpg', '2085e17b0.jpg', '803fb0bdc.jpg', \n",
    "    'f4894e99c.jpg', '193f3a261.jpg', '6b7ff665a.jpg', \n",
    "    '25563581a.jpg', 'f74fba65a.jpg', 'b5e2b8d7d.jpg'\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for cid in missing_segs:\n",
    "    if cid in train_df['ImageId'].values:\n",
    "        train_df = train_df[train_df['ImageId'] != cid]\n",
    "    else:\n",
    "        print(cid, f'not found. {cid}, (low_res)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(train_df))\n",
    "train_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df.to_csv(train_meta_file, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df.describe(percentiles=[0.01, 0.05, .25, .5, .75, 0.95, 0.99])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 64\n",
    "BATCH_NUM = 300\n",
    "\n",
    "best_left = BATCH_NUM*BATCH_SIZE\n",
    "best_right = (BATCH_NUM+1)*BATCH_SIZE\n",
    "print(best_left, best_right)\n",
    "best_idx = loss_weight_indices[best_left:best_right]\n",
    "best_loss = train_gen.loss_weights[best_idx]\n",
    "print(best_loss)\n",
    "best = train_gen.img_ids[best_idx]\n",
    "best_loss_dict = {k:v for k, v in zip(best, best_loss)}\n",
    "\n",
    "worst_left = -(BATCH_NUM+1)*BATCH_SIZE\n",
    "worst_right = -BATCH_NUM*BATCH_SIZE\n",
    "print(worst_left, worst_right)\n",
    "worst_idx = loss_weight_indices[worst_left:][::-1] if worst_right == 0 else loss_weight_indices[worst_left:worst_right][::-1]\n",
    "worst_loss = train_gen.loss_weights[worst_idx]\n",
    "print(worst_loss)\n",
    "worst = train_gen.img_ids[worst_idx]\n",
    "worst_loss_dict = {k:v for k, v in zip(worst, worst_loss)}\n",
    "\n",
    "best_worst_minmax_score = {'best': (best_loss[0], best_loss[-1]), 'worst': (worst_loss[0], worst_loss[-1])}\n",
    "loss_dict = dict(**best_loss_dict, **worst_loss_dict)\n",
    "best_worst = {'best': best, 'worst': worst}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_images = np.empty((BATCH_SIZE, CROP_SIZE, CROP_SIZE, 3), dtype=np.float32)\n",
    "samples_pseudo = np.empty((BATCH_SIZE, CROP_SIZE, CROP_SIZE, 1), dtype=np.float32)\n",
    "score_str = 'worst'\n",
    "\n",
    "for i, c_img_id in enumerate(best_worst[score_str]):\n",
    "    c_file = os.path.join(train_img_path, c_img_id)\n",
    "    c_img = cv2.cvtColor(cv2.imread(c_file), cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0\n",
    "    c_seg = seg_model.predict(c_img[None, ...])\n",
    "    samples_images[i] = c_img\n",
    "    samples_pseudo[i] = c_seg[0]\n",
    "#     samples_bboxes[i] = masks_to_bounding_boxes(c_seg)\n",
    "\n",
    "batch_rgb = montage_rgb(samples_images)\n",
    "batch_seg = montage(samples_pseudo[..., 0])\n",
    "print(samples_images.shape)\n",
    "print(batch_rgb.shape, batch_rgb.dtype)\n",
    "print(batch_seg.shape, batch_seg.dtype)\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize = (16, 16))\n",
    "ax.imshow(batch_rgb, vmin=0, vmax=1)\n",
    "ax.imshow(batch_seg, vmin=0, vmax=1, alpha=0.7)\n",
    "minmax_score = best_worst_minmax_score[score_str]\n",
    "ax.set_title(f'{score_str} (brier min:{minmax_score[0]} max:{minmax_score[-1]})')\n",
    "plt.axis('off')\n",
    "plt.savefig(os.path.join('out', BASE_MODEL, f\"{train_meta_filebase}_{score_str}_batch_{BATCH_NUM}.jpg\"))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# a ship shipping ship shipping shipping ships\n",
    "'0cb7e5f79_1.jpg'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "n_samples = 8\n",
    "aa = 0\n",
    "\n",
    "# test_files = best[aa * n_samples: (aa + 1) * n_samples]\n",
    "test_files = worst[aa * n_samples: (aa + 1) * n_samples]#[::-1]\n",
    "print(test_files)\n",
    "\n",
    "fig, m_axs = plt.subplots(n_samples, 3, figsize = (16, 42))\n",
    "with h5py.File(train_seg_file, 'r') as mask_1:\n",
    "    for (ax1, ax2, ax4), c_img_id in zip(m_axs, test_files):\n",
    "        c_file = os.path.join(train_img_path, c_img_id)\n",
    "        c_img = cv2.cvtColor(cv2.imread(c_file), cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0\n",
    "        c_seg = seg_model.predict(c_img[None, ...])\n",
    "        cur_seg = binary_opening(c_seg[0] > 0.5, np.expand_dims(disk(2), -1))\n",
    "        cur_rles = multi_rle_encode(cur_seg)\n",
    "        try:\n",
    "            y_true = mask_1[c_img_id][:] * 1\n",
    "            print('1', y_true.dtype, y_true.shape, brier_loss(y_true, c_seg))\n",
    "        except KeyError:\n",
    "            y_true = np.zeros((IMG_SIZE, IMG_SIZE, 1), dtype=np.uint8)\n",
    "            print('2', y_true.dtype, y_true.shape, brier_loss(y_true, c_seg))\n",
    "        score = seg_model.evaluate(c_img[None, ...], y_true[None, ...])\n",
    "        \n",
    "        ax1.imshow(c_img)\n",
    "        ax1.set_title(c_img_id)\n",
    "        ax2.imshow(c_seg[0, :, :, 0], vmin = 0, vmax = 1)\n",
    "        ax2.set_title('  '.join(map(lambda x: f\"{x:.3f}\", score)))\n",
    "#         ax3.imshow(cur_seg[..., 0], vmin = 0, vmax = 1)\n",
    "#         ax3.set_title(f'Pred: {len(cur_rles)}')# {loss_dict[c_img_id]}')\n",
    "#         ax3.imshow(y_true[..., -1], vmin = 0, vmax = 1)\n",
    "#         ax3.set_title(f\"Truth: \") #\"{image_weights.loc[image_weights['ImageId'] == c_img_id, 'counts'].values[0]}\")\n",
    "        ax4.imshow(c_img)\n",
    "        ax4.imshow(y_true[..., -1], vmin = 0, vmax = 1, alpha=0.3)\n",
    "#         ax4.imshow(c_seg[0, :, :, 0], vmin = 0, vmax = 1, alpha=0.3)\n",
    "        ax4.set_title(f'brier: {loss_dict[c_img_id]}')\n",
    "\n",
    "plt.tight_layout()\n",
    "fig.savefig(os.path.join('out', BASE_MODEL, f\"{train_meta_filebase}_{score_str}_batch_{BATCH_NUM}_row_{aa+1}.jpg\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predictions_to_bounding_boxes(labeled_mask):\n",
    "    if labeled_mask.max() == 0:\n",
    "        return labeled_mask\n",
    "    else:\n",
    "        img_box = np.zeros_like(labeled_mask)\n",
    "        for label_id in range(1, labeled_mask.max() + 1, 1):\n",
    "            label = np.where(labeled_mask == label_id, 1, 0).astype(np.uint8)\n",
    "            \n",
    "            _, cnt, hierarchy = cv2.findContours(label, 1, 2)\n",
    "            for c, h in zip(cnt, hierarchy[0]):\n",
    "                rect = cv2.minAreaRect(c)\n",
    "                box = cv2.boxPoints(rect)\n",
    "                box = np.int0(box)\n",
    "#                 print(c.shape, box.shape)\n",
    "                ret = cv2.matchShapes(c, box, 1, 0.0)\n",
    "#                 print(ret)\n",
    "                cv2.drawContours(img_box, [box], 0, label_id, -1)\n",
    "#                 im2, contours, hierarchy = cv2.findContours(thresh,2,1)\n",
    "#                 cnt1 = contours[0]\n",
    "#                 im2, contours, hierarchy = cv2.findContours(thresh2,2,1)\n",
    "#                 cnt2 = contours[0]\n",
    "\n",
    "    \n",
    "#             print(img_box.sum())\n",
    "        return img_box"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### [Next, Previous, First_Child, Parent]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2, 3, figsize = (16, 10))\n",
    "c_img_id = '37d94d3bc.jpg' # 1 ship, big, clear, open sea, high res, high zoom\n",
    "# c_img_id = '4830cb243.jpg' # n ship, big, clear, in port, high res, high zoom\n",
    "# c_img_id = '9b7785d98.jpg' # n ship, big, clear, in port, high res, high zoom\n",
    "# c_img_id = 'ca2a0d3ed.jpg' # n ship, big, clear, in port, high res, high zoom\n",
    "# c_img_id = '050724be2.jpg' # n ship, big, clear, in port, high res, high zoom\n",
    "# c_img_id = '00b21150c.jpg' # n ship, big, clear, in port, high res, high zoom\n",
    "# c_img_id = '8432203a3.jpg' # n ship, big, clear, in port, high res, high zoom\n",
    "# c_img_id = '2ee1dcb60.jpg' # n ships, big, clear, in port, high res, high zoom\n",
    "# c_img_id = '3926c36dd.jpg'\n",
    "# c_img_id = '34a097ff2.jpg' # n ships, big, clear, in port, high res, high zoom\n",
    "# c_img_id = '6928ac085.jpg' # n ships, big, clear, in port, high res, high zoom\n",
    "# c_img_id = '601bdf33e.jpg' # n ships, big, clear, in port, high res, high zoom\n",
    "# c_img_id = 'f6e008b23.jpg' # n ships, big, clear, in port, high res, high zoom\n",
    "# c_img_id = 'd328654ce.jpg'\n",
    "c_img_id = '4413f2ac8.jpg' # hazy\n",
    "# c_img_id = 'eba27cc8a.jpg'\n",
    "c_file = os.path.join(train_img_path, c_img_id)\n",
    "c_img = cv2.cvtColor(cv2.imread(c_file), cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0\n",
    "c_seg = seg_model.predict(c_img[None, ...])\n",
    "cur_seg = binary_opening(c_seg[0] > 0.5, np.expand_dims(disk(2), -1))\n",
    "with h5py.File(train_seg_file, 'r') as mask_1:\n",
    "    y_true = mask_1[c_img_id][:] * 1\n",
    "labeled_mask, lbl_cnt = ndi.label(cur_seg)\n",
    "# print(labeled_mask.shape, lbl_cnt, labeled_mask.min(), labeled_mask.max())\n",
    "img_box = predictions_to_bounding_boxes(labeled_mask)\n",
    "# print(img_box.shape, img_box.min(), img_box.max())\n",
    "\n",
    "img_box_flat = np.zeros(img_box.shape)\n",
    "img_box_flat[img_box > 0] = 1\n",
    "\n",
    "ax1.imshow(c_img)\n",
    "# ax1.imshow(img_box[..., 0], alpha=0.4)\n",
    "# ax1.imshow(y_true[..., -1], vmin = 0, vmax = 1, alpha=0.5)\n",
    "# ax1.imshow(labeled_mask[..., 0], alpha=0.4)\n",
    "# ax1.set_title(c_img_id)\n",
    "ax1.set_axis_off()\n",
    "\n",
    "# ax2.imshow(c_img)\n",
    "# ax2.imshow(cur_seg[..., 0], alpha=0.4)\n",
    "# ax2.imshow(c_seg[0, :, :, 0], vmin = 0, vmax = 1, alpha=0.4)\n",
    "# ax2.set_title('  '.join(map(lambda x: f\"{x:.3f}\", score)))\n",
    "\n",
    "ax2.imshow(c_seg[0, :, :, 0], vmin = 0, vmax = 1, alpha=1.0)\n",
    "ax2.set_axis_off()\n",
    "ax3.imshow(cur_seg[..., 0], alpha=1.0)\n",
    "ax3.set_axis_off()\n",
    "ax4.imshow(labeled_mask[..., 0], alpha=1.0)\n",
    "ax4.set_axis_off()\n",
    "ax5.imshow(img_box[..., 0], alpha=1.0)\n",
    "ax5.set_axis_off()\n",
    "ax6.imshow(c_img)\n",
    "ax6.imshow(img_box_flat[..., 0], alpha=0.4)\n",
    "ax6.set_axis_off()\n",
    "# ax6.imshow(y_true[..., -1], vmin = 0, vmax = 1, alpha=1.0)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize = (12, 4))\n",
    "# c_img_id = '37d94d3bc.jpg' # 1 ship, big, clear, open sea, high res, high zoom\n",
    "# c_img_id = '4830cb243.jpg' # n ship, big, clear, in port, high res, high zoom\n",
    "# c_img_id = '9b7785d98.jpg' # n ship, big, clear, in port, high res, high zoom\n",
    "# c_img_id = 'ca2a0d3ed.jpg' # n ship, big, clear, in port, high res, high zoom\n",
    "# c_img_id = '050724be2.jpg' # n ship, big, clear, in port, high res, high zoom\n",
    "# c_img_id = '00b21150c.jpg' # n ship, big, clear, in port, high res, high zoom\n",
    "# c_img_id = '8432203a3.jpg' # n ship, big, clear, in port, high res, high zoom\n",
    "# c_img_id = '2ee1dcb60.jpg' # n ships, big, clear, in port, high res, high zoom\n",
    "# c_img_id = '601bdf33e.jpg' # n ships, big, clear, in port, high res, high zoom\n",
    "# c_img_id = 'f6e008b23.jpg' # n ships, big, clear, in port, high res, high zoom\n",
    "# c_img_id = 'eba27cc8a.jpg'\n",
    "c_img_id = '3926c36dd.jpg'\n",
    "c_img_id = '34a097ff2.jpg' # n ships, big, clear, in port, high res, high zoom\n",
    "c_img_id = '4413f2ac8.jpg' # hazy\n",
    "c_img_id = 'd328654ce.jpg'\n",
    "c_img_id = '6928ac085.jpg' # n ships, big, clear, in port, high res, high zoom\n",
    "c_file = os.path.join(train_img_path, c_img_id)\n",
    "c_img = cv2.cvtColor(cv2.imread(c_file), cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0\n",
    "c_seg = seg_model.predict(c_img[None, ...])\n",
    "cur_seg = binary_opening(c_seg[0] > 0.5, np.expand_dims(disk(2), -1))\n",
    "with h5py.File(train_seg_file, 'r') as mask_1:\n",
    "    y_true = mask_1[c_img_id][:] * 1\n",
    "labeled_mask, lbl_cnt = ndi.label(cur_seg)\n",
    "img_box = predictions_to_bounding_boxes(labeled_mask)\n",
    "\n",
    "img_box_flat = np.zeros(img_box.shape)\n",
    "img_box_flat[img_box > 0] = 1\n",
    "\n",
    "ax1.imshow(c_img)\n",
    "# ax1.imshow(img_box_flat[..., 0], alpha=1.0)\n",
    "ax1.set_title('input')\n",
    "ax1.set_axis_off()\n",
    "\n",
    "# ax1.imshow(c_img)\n",
    "ax2.imshow(img_box_flat[..., 0], alpha=1.0)\n",
    "ax2.set_title('prediction')\n",
    "ax2.set_axis_off()\n",
    "\n",
    "ax3.imshow(c_img)\n",
    "ax3.imshow(img_box_flat[..., 0], alpha=0.4)\n",
    "ax3.set_title('overlay')\n",
    "ax3.set_axis_off()\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plt.tight_layout()\n",
    "fig.savefig(os.path.join('out', BASE_MODEL, f\"3_{c_img_id}\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fig, (ax1, ax2, ax3, ax4, ax5) = plt.subplots(1, 5, figsize = (14, 3))\n",
    "c_img_id = 'eba27cc8a.jpg'\n",
    "c_file = os.path.join(train_img_path, c_img_id)\n",
    "c_img = cv2.cvtColor(cv2.imread(c_file), cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0\n",
    "c_seg = seg_model.predict(c_img[None, ...])\n",
    "cur_seg = binary_opening(c_seg[0] > 0.5, np.expand_dims(disk(2), -1))\n",
    "with h5py.File(train_seg_file, 'r') as mask_1:\n",
    "    y_true = mask_1[c_img_id][:] * 1\n",
    "labeled_mask, lbl_cnt = ndi.label(cur_seg)\n",
    "# print(labeled_mask.shape, lbl_cnt, labeled_mask.min(), labeled_mask.max())\n",
    "img_box = predictions_to_bounding_boxes(labeled_mask)\n",
    "# print(img_box.shape, img_box.min(), img_box.max())\n",
    "\n",
    "ax1.imshow(c_img)\n",
    "# ax1.imshow(img_box[..., 0], alpha=0.4)\n",
    "# ax1.imshow(y_true[..., -1], vmin = 0, vmax = 1, alpha=0.5)\n",
    "# ax1.imshow(labeled_mask[..., 0], alpha=0.4)\n",
    "# ax1.set_title(c_img_id)\n",
    "ax1.set_title('input to NN')\n",
    "ax1.set_axis_off()\n",
    "\n",
    "# ax2.imshow(c_img)\n",
    "# ax2.imshow(cur_seg[..., 0], alpha=0.4)\n",
    "# ax2.imshow(c_seg[0, :, :, 0], vmin = 0, vmax = 1, alpha=0.4)\n",
    "# ax2.set_title('  '.join(map(lambda x: f\"{x:.3f}\", score)))\n",
    "\n",
    "ax2.imshow(c_seg[0, :, :, 0], vmin = 0, vmax = 1, alpha=1.0)\n",
    "ax2.set_title('NN output')\n",
    "ax2.set_axis_off()\n",
    "ax3.imshow(cur_seg[..., 0], alpha=1.0)\n",
    "ax3.set_title('thresholding')\n",
    "ax3.set_axis_off()\n",
    "ax4.imshow(img_box[..., 0], alpha=1.0)\n",
    "ax4.set_title('bbox')\n",
    "ax4.set_axis_off()\n",
    "ax5.imshow(c_img)\n",
    "ax5.imshow(img_box[..., 0], alpha=0.4)\n",
    "ax5.set_title('final overlay')\n",
    "ax5.set_axis_off()\n",
    "# ax6.imshow(y_true[..., -1], vmin = 0, vmax = 1, alpha=1.0)\n",
    "\n",
    "# plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plt.tight_layout()\n",
    "fig.savefig(os.path.join('out', BASE_MODEL, f\"5_{c_img_id}\"))"
   ]
  }
 ],
 "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.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}