dataset/b-it-bots_dataset_preprocessing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# b-it-bots@Work dataset preprocessing\n",
    "\n",
    "* The dataset directory structuce must follow:\n",
    "  ```\n",
    "  b-it-bots_atwork_dataset\n",
    "  ├── train\n",
    "  │   ├── AXIS\n",
    "  │   ├── ...\n",
    "  ├── test\n",
    "  │   ├── AXIS\n",
    "  │   ├── ...\n",
    "  ```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import open3d\n",
    "import numpy as np\n",
    "import h5py\n",
    "import pickle\n",
    "import sklearn\n",
    "from sklearn.decomposition import PCA\n",
    "import os\n",
    "import glob\n",
    "import dataset_util"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load data open3d==0.10.0\n",
    "#dataset_dir = '/media/mhwasil/WDEXT11/HDD4/pointcloud_dataset/pcd/b-it-bots/combined_pointcloud_dataset'\n",
    "dataset_dir = \"/media/mhwasil/WDEXT11/HDD4/pointcloud_dataset/pcd/b-it-bots/realsense_d435_2021\"\n",
    "dataset_name = \"b_it_bots_at_work_2021\"\n",
    "train_dir = os.path.join(dataset_dir, \"train\")\n",
    "test_dir = os.path.join(dataset_dir, \"test\")\n",
    "all_dataset = [train_dir, test_dir]\n",
    "\n",
    "# downsample and padding param\n",
    "# downsample if > 2048, and pad if < 2048\n",
    "num_points = 2048\n",
    "\n",
    "# split train or test files into multiple files\n",
    "number_train_file = 4\n",
    "number_test_file = 1\n",
    "\n",
    "total_dataset = 0\n",
    "\n",
    "for split in [\"train\", \"test\"]:\n",
    "    data_path = os.path.join(dataset_dir, split)\n",
    "    train_classes = [label for label in os.listdir(data_path) if os.path.isdir(os.path.join(data_path, label))]\n",
    "    print(\"Dataset: \", split)\n",
    "    data = []\n",
    "    labels = []\n",
    "    for i,label in enumerate (train_classes):\n",
    "        pcd_files = os.listdir(os.path.join(data_path, label))\n",
    "        np.random.shuffle(pcd_files)\n",
    "        print (\"Preprocessing \",label)\n",
    "        data_per_class = []\n",
    "        total_per_class = 0\n",
    "        for j, pcd_file_name in enumerate(pcd_files):\n",
    "            pcd_path = os.path.join(data_path, label, pcd_file_name)\n",
    "            xyzrgb = dataset_util.extract_pcd(pcd_path, num_points=num_points, \n",
    "                                      color=True, downsample_cloud=True, \n",
    "                                      pad_cloud=True, normalize_cloud=True)\n",
    "            \n",
    "            if xyzrgb is not None:\n",
    "                labels.append(i)\n",
    "                data.append(xyzrgb)\n",
    "\n",
    "                total_dataset += 1\n",
    "                total_per_class += 1\n",
    "            \n",
    "    data = np.asarray(data)\n",
    "    labels = np.asarray(labels)\n",
    "    print (\"Data\", data.shape)\n",
    "    print (\"Labels\", labels.shape)\n",
    "    \n",
    "    #shuffle dataset\n",
    "    data, labels = dataset_util.randomize(data, labels)\n",
    "    \n",
    "    \n",
    "    # data perfile\n",
    "    # todo: split into multiple files for large dataset\n",
    "    #data_perfile = int(data.shape[0]/number_train_file)\n",
    "    #print(data_perfile)\n",
    "    \n",
    "    data_dict = {}\n",
    "    data_dict['data'] = data\n",
    "    data_dict['labels'] = labels\n",
    "    \n",
    "    # save pickle file\n",
    "    dataset_util.save_dataset_and_compress(data_dict, dataset_dir+'/{}_{}'.format(dataset_name, split))\n",
    "    \n",
    "    data = []\n",
    "    labels = []\n",
    "    data_dict.clear()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:deeplearning] *",
   "language": "python",
   "name": "conda-env-deeplearning-py"
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}