prepare_date.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "secure-windows",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "import cv2 \n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "%pylab inline\n",
    "%config Completer.use_jedi = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "constitutional-practice",
   "metadata": {},
   "outputs": [],
   "source": [
    "SEED = 42\n",
    "\n",
    "BBOX_MARGIN = .2\n",
    "TIME_STAMP = '2021_05_03'\n",
    "\n",
    "ITASCA_PATH = '/home/ubuntu/percepto/data/ItascaClassification'\n",
    "\n",
    "SRC_PATH = os.path.join(ITASCA_PATH, 'in_progress')\n",
    "\n",
    "TARGET_PATH = os.path.join(ITASCA_PATH, 'dataset')\n",
    "IMAGE_PATH = os.path.join(TARGET_PATH, 'images')\n",
    "ANNOTATIONS_PATH = os.path.join(TARGET_PATH, 'annotations')\n",
    "MAPPING_PATH = os.path.join(TARGET_PATH, 'mapping', TIME_STAMP)\n",
    "\n",
    "MAPPING_DICT = {'ignore': 0,\n",
    "                None: 2000, \n",
    "                'semantic_change': 2001,\n",
    "                'air-leak': 2002, \n",
    "                'ground-leak': 2003, \n",
    "                'fire': 2004, \n",
    "                'gnome': 2005, \n",
    "                'leak': 2006}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "crude-edwards",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! mkdir {os.path.join(ITASCA_PATH, 'dataset', 'appendix')}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "palestinian-disabled",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(ITASCA_PATH, 'dataset', 'appendix', 'class_mapping.json'), 'w') as fp:\n",
    "    json.dump(MAPPING_DICT, fp)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dressed-chuck",
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted(os.listdir(SRC_PATH))[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "scheduled-fiction",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(list(filter(lambda s: s.endswith('jpg'), os.listdir(SRC_PATH))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "massive-financing",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 0\n",
    "\n",
    "class_set = set()\n",
    "\n",
    "for file in os.listdir(SRC_PATH):\n",
    "    \n",
    "    if file.endswith('jpg'):\n",
    "    \n",
    "        image_file = os.path.join(SRC_PATH, file)\n",
    "        json_file = os.path.join(SRC_PATH, f'{file}___pixel.json')\n",
    "        mask_file = os.path.join(SRC_PATH, f'{file}___fuse.png')\n",
    "\n",
    "        with open(json_file) as jf:\n",
    "            data = json.load(jf)\n",
    "        if len(data['instances']) and (data['instances'][0]['className'] != 'semantic_change'):\n",
    "            i += 1\n",
    "\n",
    "            class_name = data['instances'][0]['className']\n",
    "\n",
    "            if class_name not in class_set:\n",
    "                class_set.add(class_name)\n",
    "                image = cv2.imread(image_file)\n",
    "                mask = cv2.imread(mask_file)\n",
    "\n",
    "                plt.figure(figsize=(20, 10))\n",
    "                plt.imshow(np.hstack([image[...,::-1], mask]))\n",
    "                plt.title([i, data['instances'][0]['className']])\n",
    "print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "parental-necessity",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 0\n",
    "for file in os.listdir(SRC_PATH):\n",
    "    \n",
    "    if file.endswith('jpg'):\n",
    "\n",
    "        image_file = os.path.join(SRC_PATH, file)\n",
    "        json_file = os.path.join(SRC_PATH, f'{file}___pixel.json')\n",
    "        mask_file = os.path.join(SRC_PATH, f'{file}___fuse.png')\n",
    "\n",
    "        with open(json_file) as jf:\n",
    "            data = json.load(jf)\n",
    "\n",
    "        if len(data['instances']) and (data['instances'][0]['className'] == 'ground-leak'):\n",
    "            i += 1\n",
    "            class_name = 'ground-leak'\n",
    "\n",
    "            image = cv2.imread(image_file)\n",
    "            mask = cv2.imread(mask_file)\n",
    "\n",
    "            plt.figure(figsize=(20, 10))\n",
    "            plt.imshow(np.hstack([image[...,::-1], mask]))\n",
    "            plt.title([i, data['instances'][0]['className'], file[:-4]])\n",
    "\n",
    "print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "accomplished-smell",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 0\n",
    "for file in os.listdir(SRC_PATH):\n",
    "    \n",
    "    if file.endswith('jpg'):\n",
    "\n",
    "        image_file = os.path.join(SRC_PATH, file)\n",
    "        json_file = os.path.join(SRC_PATH, f'{file}___pixel.json')\n",
    "        mask_file = os.path.join(SRC_PATH, f'{file}___fuse.png')\n",
    "    \n",
    "        with open(json_file) as jf:\n",
    "            data = json.load(jf)\n",
    "\n",
    "        if len(data['instances']) and (data['instances'][0]['className'] == 'air-leak'):\n",
    "            i += 1\n",
    "            class_name = 'air-leak'\n",
    "\n",
    "            image = cv2.imread(image_file)\n",
    "            mask = cv2.imread(mask_file)\n",
    "\n",
    "            plt.figure(figsize=(20, 10))\n",
    "            plt.imshow(np.hstack([image[...,::-1], mask]))\n",
    "            plt.title([i, data['instances'][0]['className'], file[:-4]])\n",
    "        \n",
    "print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sophisticated-objective",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 0\n",
    "\n",
    "for file in os.listdir(SRC_PATH):\n",
    "    \n",
    "    if file.endswith('jpg'):\n",
    "\n",
    "        image_file = os.path.join(SRC_PATH, file)\n",
    "        json_file = os.path.join(SRC_PATH, f'{file}___pixel.json')\n",
    "        mask_file = os.path.join(SRC_PATH, f'{file}___fuse.png')\n",
    "    \n",
    "    with open(json_file) as jf:\n",
    "        data = json.load(jf)\n",
    "        \n",
    "        if len(data['instances']) and (data['instances'][0]['className'] != 'semantic_change'):\n",
    "            i += 1\n",
    "\n",
    "            class_name = data['instances'][0]['className'].split('-')[-1]\n",
    "\n",
    "            image = cv2.imread(image_file)\n",
    "            mask = cv2.imread(mask_file)\n",
    "            maskgray = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)\n",
    "            _, thresh = cv2.threshold(maskgray,127,255,2)\n",
    "            contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n",
    "\n",
    "            max_area = 0 \n",
    "            for c in contours:\n",
    "                area = cv2.contourArea(c)\n",
    "                if area > max_area:\n",
    "                    cnt = c\n",
    "                    max_area = area\n",
    "\n",
    "            x, y, w, h = cv2.boundingRect(cnt)\n",
    "            dw = int(BBOX_MARGIN*w)\n",
    "            dh = int(BBOX_MARGIN*h)\n",
    "            cv2.rectangle(image,(x,y),(x+w,y+h),(0,255,0),5)\n",
    "            cv2.rectangle(image,(x-dw,y-dh),(x+w+dw,y+h+dh),(0,0,255),5)\n",
    "\n",
    "            plt.figure(figsize=(20, 15))\n",
    "            plt.imshow(np.hstack([image[...,::-1], mask]))\n",
    "\n",
    "            plt.title([i, class_name, len(contours)])\n",
    "\n",
    "            if i == 10:\n",
    "                break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "unavailable-cornell",
   "metadata": {},
   "outputs": [],
   "source": [
    "def crop_out_of_bounds(img, bbox):\n",
    "    \"\"\"crops an image when bounding box can be out of image bounderies\"\"\"\n",
    "    \n",
    "    h, w = img.shape[:2]\n",
    "    x, y, dx, dy = bbox\n",
    "    \n",
    "    base_crop = img[max(0, y):y+dy, max(0, x):x+dx].copy()\n",
    "    padded_crop = np.zeros((dy, dx, 3), dtype='uint8')\n",
    "    padded_crop[abs(min(0, y)):h-y, abs(min(0, x)):w-x] = base_crop\n",
    "    \n",
    "    return padded_crop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "extra-equation",
   "metadata": {},
   "outputs": [],
   "source": [
    "files_list = []\n",
    "mapping_list = []\n",
    "classes_list = []\n",
    "bbox_list = []\n",
    "\n",
    "i = 0\n",
    "\n",
    "crop_ext = f'_cropped_{BBOX_MARGIN}.jpg'\n",
    "\n",
    "for file in os.listdir(SRC_PATH):\n",
    "    \n",
    "    if file.endswith('jpg'):\n",
    "\n",
    "        image_file = os.path.join(SRC_PATH, file)\n",
    "        json_file = os.path.join(SRC_PATH, f'{file}___pixel.json')\n",
    "        mask_file = os.path.join(SRC_PATH, f'{file}___fuse.png')\n",
    "    \n",
    "        with open(json_file) as jf:\n",
    "            data = json.load(jf)\n",
    "\n",
    "        if len(data['instances']) and (data['instances'][0]['className'] != 'semantic_change'):\n",
    "            i += 1\n",
    "\n",
    "            class_name = data['instances'][0]['className'].split('-')[-1]\n",
    "\n",
    "            image = cv2.imread(image_file)\n",
    "            mask = cv2.imread(mask_file)\n",
    "            mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)\n",
    "            _, thresh = cv2.threshold(mask,127,255,2)\n",
    "            contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n",
    "\n",
    "            max_area = 0 \n",
    "            for c in contours:\n",
    "                area = cv2.contourArea(c)\n",
    "                if area > max_area:\n",
    "                    cnt = c\n",
    "                    max_area = area\n",
    "\n",
    "            x, y, w, h = cv2.boundingRect(cnt)\n",
    "\n",
    "            dw = int(BBOX_MARGIN*w)\n",
    "            dh = int(BBOX_MARGIN*h)\n",
    "\n",
    "            cropped_img = crop_out_of_bounds(image, (x-dw, y-dh, w+2*dw, h+2*dh))\n",
    "\n",
    "            cropped_file = file.replace('.jpg', crop_ext)\n",
    "            annotation_file = cropped_file.replace(\"jpg\", \"txt\")\n",
    "\n",
    "            cv2.imwrite(os.path.join(IMAGE_PATH, cropped_file), cropped_img)\n",
    "\n",
    "            top, left, bottom, right = dh, dw, dh+h, dw+w\n",
    "\n",
    "            with open(os.path.join(ANNOTATIONS_PATH, annotation_file), 'w') as f:\n",
    "                      f.write(f'{MAPPING_DICT[class_name]},1,0,0,{left},{top},{right},{top},{right},{bottom},{left},{bottom}')\n",
    "\n",
    "            files_list.append(f'images/{cropped_file}')\n",
    "            mapping_list.append(f'annotations/{annotation_file}')\n",
    "            classes_list.append(MAPPING_DICT[class_name])\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "conceptual-vegetarian",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "df = pd.DataFrame(list(zip(files_list, mapping_list, classes_list)), \n",
    "                  columns =['image_name', 'mapping_file', 'class'])\n",
    "\n",
    "print(df.head())\n",
    "                \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "actual-melissa",
   "metadata": {},
   "outputs": [],
   "source": [
    "list(df.image_name[:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "rising-resistance",
   "metadata": {},
   "outputs": [],
   "source": [
    "list(df.mapping_file[:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "functioning-relative",
   "metadata": {},
   "outputs": [],
   "source": [
    "df['class'].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dying-ultimate",
   "metadata": {},
   "outputs": [],
   "source": [
    "train, val = train_test_split(df, random_state=SEED, stratify=df['class'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "monthly-cartridge",
   "metadata": {},
   "outputs": [],
   "source": [
    "train['class'].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dried-croatia",
   "metadata": {},
   "outputs": [],
   "source": [
    "val['class'].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "harmful-activity",
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir -p {MAPPING_PATH}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "characteristic-terminology",
   "metadata": {},
   "outputs": [],
   "source": [
    "train.to_csv(os.path.join(MAPPING_PATH, 'train.txt'), columns=['image_name', 'mapping_file'], index=False, header=False)\n",
    "val.to_csv(os.path.join(MAPPING_PATH, 'val.txt'), columns=['image_name', 'mapping_file'], index=False, header=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecological-valve",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_animation(path, relative_path=''):\n",
    "    \n",
    "    animation_path = os.path.join(ITASCA_PATH, 'animation', os.path.splitext(os.path.basename(path))[0])\n",
    "    print(animation_path)\n",
    "\n",
    "    with open(path, 'r') as f:\n",
    "        files_list = f.readlines()\n",
    "    \n",
    "    for i, line in enumerate(files_list):\n",
    "        image_path, annotation_path = line.strip().split(',')\n",
    "        image_path = os.path.join(relative_path,'dataset', image_path)\n",
    "        annotation_path = os.path.join(relative_path,'dataset', annotation_path)\n",
    "        \n",
    "        image = cv2.imread(image_path)\n",
    "        \n",
    "        with open(annotation_path, 'r') as f:\n",
    "            class_label, _, _, _, left, top, right, _, _, bottom, _, _ = [int(x) for x in f.read().split(',')]\n",
    "        \n",
    "        if class_label == 2004:\n",
    "            c = (255, 0, 0)\n",
    "        elif class_label == 2005:\n",
    "            c = (0, 255, 0)\n",
    "        else:\n",
    "            c = (0, 0, 255)\n",
    "        \n",
    "        image_w_bbox = image.copy()\n",
    "\n",
    "        cv2.rectangle(image_w_bbox,(left, top), (right, bottom), c, -1)\n",
    "        alpha = 0.3\n",
    "        cv2.addWeighted(image_w_bbox, alpha, image, 1 - alpha, 0, image_w_bbox)\n",
    "                \n",
    "        name_to_save = os.path.join(animation_path, f'{i:04d}_*.jpg')\n",
    "        cv2.imwrite(name_to_save.replace('*', 'image'), image)\n",
    "        cv2.imwrite(name_to_save.replace('*', 'w_bbox'), image_w_bbox)\n",
    "                                            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "concrete-editing",
   "metadata": {},
   "outputs": [],
   "source": [
    "make_animation(os.path.join(MAPPING_PATH, 'train.txt'), ITASCA_PATH)\n",
    "make_animation(os.path.join(MAPPING_PATH, 'val.txt'), ITASCA_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "warming-emperor",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}