FileIO.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run Utility_general.ipynb\n",
    "import glob\n",
    "import struct\n",
    "from pathlib import Path\n",
    "from PIL import Image as PImg\n",
    "from shutil import copyfile\n",
    "import cv2\n",
    "import numpy as np\n",
    "import nibabel as nib\n",
    "import os\n",
    "#import h5py\n",
    "\n",
    "\n",
    "from matplotlib import image as mpimg\n",
    "\n",
    "# Refer to https://docs.python.org/3/library/struct.html\n",
    "# f: float; d: double; q: long long; l: long; i: int; I: unsigned int\n",
    "dlength_dict = {'f': 4, 'd': 8, 'q': 8, 'l': 4, 'i': 4, 'I': 4}\n",
    "\n",
    "class FileIO_DMT(object):\n",
    "    \n",
    "    DIPHA_CONST = 8067171840\n",
    "    DIPHA_IMAGE_TYPE_CONST = 1\n",
    "    \n",
    "    @staticmethod\n",
    "    def cvt_dat2D2dipha_(dat, pathOut):\n",
    "        '''\n",
    "        img: the data to be written out as dipha\n",
    "        pathOut: path to the binary output file\n",
    "        output file format compatible with DMT program\n",
    "        output file name has to end with .bin\n",
    "        '''\n",
    "        assert(pathOut.endswith('.bin'))\n",
    "        assert(size(dat.shape) == 2)\n",
    "        fout = open(pathOut, \"wb\")\n",
    "        fout.write(struct.pack('q', FileIO_DMT.DIPHA_CONST))\n",
    "        fout.write(struct.pack('q', FileIO_DMT.DIPHA_IMAGE_TYPE_CONST))\n",
    "        fout.write(struct.pack('q', np.prod(dat.shape)))\n",
    "        fout.write(struct.pack('q', 3))\n",
    "        fout.write(struct.pack('q', dat.shape[0]))\n",
    "        fout.write(struct.pack('q', dat.shape[1]))\n",
    "        fout.write(struct.pack('q', 1))\n",
    "        dat = np.transpose(-dat)\n",
    "        fout.write(struct.pack(str(np.prod(dat.shape))+'d', *(dat.flatten())))\n",
    "        fout.close()\n",
    "    \n",
    "    @staticmethod\n",
    "    def cvt_img2dipha_(pathIn, pathOut, transform_argv=[]):\n",
    "        '''\n",
    "        pathIn: path to the image\n",
    "        pathOut: path to the binary output file\n",
    "        - transform_argv: transform arguments. \n",
    "        - Example: ['dist_trfm', True, True]\n",
    "        '''\n",
    "        img  = cv2.imread(pathIn, cv2.IMREAD_GRAYSCALE)\n",
    "        img  = np.float64(img)\n",
    "        if len(transform_argv) != 0 and transform_argv[0] == 'dist_trfm':\n",
    "            assert(len(transform_argv) == 3)\n",
    "            dt = Util_gen.dist_trfm(img, binarize=transform_argv[1], inverse=transform_argv[2])\n",
    "            FileIO_DMT.cvt_dat2D2dipha_(dt, pathOut)\n",
    "        else:\n",
    "            FileIO_DMT.cvt_dat2D2dipha_(img, pathOut)\n",
    "        \n",
    "    @staticmethod\n",
    "    def cvt_dipha2img_(pathIn, uint8=False):\n",
    "        '''\n",
    "        pathIn: path to the binary dipha file\n",
    "        uint8: convert output matrix to uint8 data type\n",
    "        '''\n",
    "        fin = open(pathIn, \"rb\")\n",
    "        assert(struct.unpack('q', fin.read(dlength_dict['q']))[0] == FileIO_DMT.DIPHA_CONST)\n",
    "        assert(struct.unpack('q', fin.read(dlength_dict['q']))[0] == FileIO_DMT.DIPHA_IMAGE_TYPE_CONST)\n",
    "        total_pixel = struct.unpack('q', fin.read(dlength_dict['q']))[0]\n",
    "        struct.unpack('q', fin.read(dlength_dict['q']))\n",
    "        [nx, ny, nz] = struct.unpack('3q', fin.read(dlength_dict['q']*3))\n",
    "        assert(nx*ny*nz == total_pixel)\n",
    "        img = struct.unpack(str(total_pixel)+'d', fin.read(dlength_dict['d']*total_pixel))\n",
    "        img = np.reshape(img, [nx, ny])\n",
    "        img = np.transpose(-img)\n",
    "        if uint8:\n",
    "            img = np.uint8(img)\n",
    "        return img\n",
    "    \n",
    "    @staticmethod\n",
    "    def cvt_img2dipha_batch_(dirIn, dirOut, ext, transform_argv=[]):\n",
    "        '''\n",
    "        dirIn: the folder containing the images\n",
    "        dirOut: the folder to output .bin dipha files\n",
    "        ext: target extension of the images\n",
    "        transform_argv: transform arguments.\n",
    "        '''\n",
    "        filesIn  = []\n",
    "        filesOut = []\n",
    "        os.chdir(dirIn)\n",
    "        for file in glob.glob(\"*.\"+ext):\n",
    "            name = file[:-len(ext)-1]\n",
    "            filesIn.append(os.path.join(dirIn, file))\n",
    "            filesOut.append(os.path.join(dirOut, name+\".bin\"))\n",
    "        Path(dirOut).mkdir(parents=True, exist_ok=True)\n",
    "        \n",
    "        for i in range(len(filesIn)):\n",
    "            FileIO_DMT.cvt_img2dipha_(filesIn[i], filesOut[i], transform_argv)\n",
    "    \n",
    "class FileIO(object):\n",
    "    \n",
    "    @staticmethod\n",
    "    def read_binary(path, shape, dtype='d'):\n",
    "        file_in  = open(path, \"rb\")\n",
    "        data_arr = struct.unpack(str(np.prod(shape))+dtype, file_in.read(dlength_dict[dtype]*np.prod(shape)))\n",
    "        data_arr = np.reshape(data_arr, shape)\n",
    "        file_in.close()\n",
    "        if dtype=='f':\n",
    "            data_arr = np.float32(data_arr)\n",
    "        return data_arr\n",
    "    \n",
    "    @staticmethod\n",
    "    def write_binary(path, data, shape, dtype='d'):\n",
    "        '''\n",
    "        data has to be flattened.\n",
    "        '''\n",
    "        file_out = open(path, \"wb\")\n",
    "        file_out.write(struct.pack(str(np.prod(shape))+dtype, *(data)))\n",
    "        file_out.close()\n",
    "    \n",
    "    @staticmethod\n",
    "    def read_matrix_binary(path, dtype='d'):\n",
    "        file_in = open(path, \"rb\")\n",
    "        dims    = struct.unpack('I', file_in.read(4))[0]\n",
    "        if dims == 0:\n",
    "            return None\n",
    "        shape   = struct.unpack(str(dims)+'I', file_in.read(4 * dims))\n",
    "        mat     = struct.unpack(str(np.prod(shape))+dtype, file_in.read(dlength_dict[dtype]*np.prod(shape)))\n",
    "        mat     = np.reshape(mat, shape)\n",
    "        return mat\n",
    "    \n",
    "    @staticmethod\n",
    "    def write_matrix_binary(path, mat, dtype='d'):\n",
    "        dims  = len(mat.shape)\n",
    "        shape = mat.shape\n",
    "        file_out = open(path, \"wb\")\n",
    "        file_out.write(struct.pack('I', dims))\n",
    "        file_out.write(struct.pack(str(dims)+'I', *(shape)))\n",
    "        file_out.write(struct.pack(str(np.prod(shape))+dtype, *(mat.flatten())))\n",
    "        file_out.close()\n",
    "    \n",
    "    @staticmethod\n",
    "    def load_nii_(path, verbose=False):\n",
    "        vol_struct = nib.load(path)\n",
    "        vol = np.squeeze(np.array(vol_struct.get_fdata()))\n",
    "        assert(size(vol.shape) == 3)\n",
    "        print(vol.shape)\n",
    "        print(type(vol[0,0,0]))\n",
    "        \n",
    "        if verbose == True:\n",
    "            print(vol_struct.affine)\n",
    "            print(vol_struct.header)\n",
    "        return vol\n",
    "    \n",
    "#     @staticmethod\n",
    "#     def read_h5py(path, shape):\n",
    "#         with h5py.File(path, \"r\") as hf:    \n",
    "\n",
    "#             # Split the data into training/test features/targets\n",
    "#             X_train = hf[\"X_train\"][:]\n",
    "#             targets_train = hf[\"y_train\"][:]\n",
    "#             X_test = hf[\"X_test\"][:] \n",
    "#             targets_test = hf[\"y_test\"][:]\n",
    "\n",
    "#             # Determine sample shape\n",
    "#             sample_shape = shape\n",
    "\n",
    "#             # Reshape data into 3D format\n",
    "#             X_train = Util_gen.rgb_data_transform(X_train)\n",
    "#             X_test = Util_gen.rgb_data_transform(X_test)\n",
    "#         return X_train, X_test, targets_train, targets_test\n",
    "    \n",
    "    @staticmethod\n",
    "    def save_image_batch(data, folder, prefix, scalor, number_offset, binary_out, num=-1):\n",
    "        '''\n",
    "        data: the image data to be saved\n",
    "        folder: the path to the folder where the images are saved\n",
    "        prefix: prefix_00000.png\n",
    "        scalor: data * scalor + scalor, 0 indicates no scaling\n",
    "        number_offset: the index of the image to start saving from (included)\n",
    "        num: number of images to save out, -1 means all\n",
    "        '''\n",
    "        data = np.squeeze(data)\n",
    "        assert(len(data.shape) == 3 or len(data.shape) == 4)\n",
    "        batch_size  = data.shape[0]\n",
    "        image_shape = data.shape[-2:]\n",
    "        \n",
    "        if num == -1 or num >= batch_size:\n",
    "            out_num = batch_size\n",
    "        else:\n",
    "            out_num = num\n",
    "            \n",
    "        for idx in range(out_num):\n",
    "            out_name = ''\n",
    "            for _ in range(5 - len(str(idx + number_offset))):\n",
    "                out_name += '0'\n",
    "            dat_ = np.reshape(data[idx], image_shape)\n",
    "            if scalor > 0:\n",
    "                dat_ = dat_ * scalor + scalor\n",
    "            if binary_out == False:\n",
    "                out_name = folder + \"/\" + prefix + \"_\" + out_name + str(idx + number_offset) + \".png\"\n",
    "                dat_ = dat_.astype(np.uint8)\n",
    "                cv2.imwrite(out_name, dat_)\n",
    "            else:\n",
    "                out_name = folder + \"/\" + prefix + \"_\" + out_name + str(idx + number_offset) + \".dat\"\n",
    "                FileIO.write_matrix_binary(out_name, dat_, 'f')\n",
    "        print(\"Data write out complete.\")\n",
    "    \n",
    "    @staticmethod\n",
    "    def img_boundbox_batch_(dirIn, dirOut, ext, thick):\n",
    "        '''\n",
    "        dirIn: the folder containing the images\n",
    "        dirOut: the folder to output .bin dipha files\n",
    "        ext: target extension of the images\n",
    "        thick: thickness of the bounding box\n",
    "        This function reads and output images in GRAYSCALE\n",
    "        '''\n",
    "        os.chdir(dirIn)\n",
    "        Path(dirOut).mkdir(parents=True, exist_ok=True)\n",
    "        for file in glob.glob(\"*.\"+ext):\n",
    "            pathIn = os.path.join(dirIn, file)\n",
    "            pathOut = os.path.join(dirOut, file)\n",
    "            img = cv2.imread(pathIn, cv2.IMREAD_GRAYSCALE)\n",
    "            cv2.rectangle(img, (0,0),(img.shape[1]-1, img.shape[0]-1),(0), thick)\n",
    "            cv2.imwrite(pathOut, img)\n",
    "            \n",
    "    @staticmethod\n",
    "    def img_connected_component_cv_batch_(dirIn, dirOut, ext, binarize, fillRift, shuffle_num):\n",
    "        '''\n",
    "        dirIn: the folder containing the images\n",
    "        dirOut: the folder to output .bin dipha files\n",
    "        ext: target extension of the images\n",
    "        This function reads images and outputs their connected components ranging from 0 to N\n",
    "        '''\n",
    "        filesIn  = []\n",
    "        filesOut = []\n",
    "        os.chdir(dirIn)\n",
    "        Path(dirOut).mkdir(parents=True, exist_ok=True)\n",
    "        for file in glob.glob(\"*.\"+ext):\n",
    "            name = file[:-len(ext)-1]\n",
    "            filesIn.append(os.path.join(dirIn, file))\n",
    "            filesOut.append(os.path.join(dirOut, name+\".dat\"))\n",
    "        for i in range(len(filesIn)):\n",
    "            img = cv2.imread(filesIn[i], cv2.IMREAD_GRAYSCALE)\n",
    "            if binarize == True:\n",
    "                img[img>127]  = 255\n",
    "                img[img<=127] = 0\n",
    "            cnt, hry, red = Util_cv.compute_bnd_red_cv(img, 0, 255, 8)\n",
    "            if fillRift == True:\n",
    "                filled = Util_gen.fill_rift_(red[1], 8)\n",
    "                if shuffle_num > 0:\n",
    "                    Util_gen.shuffle_partition_label(filled, shuffle_num, filesOut[i])\n",
    "                else:\n",
    "                    FileIO.write_matrix_binary(filesOut[i], filled, 'i')\n",
    "            else:\n",
    "                if shuffle_num > 0:\n",
    "                    Util_gen.shuffle_partition_label(red[1], shuffle_num, filesOut[i])\n",
    "                else:\n",
    "                    FileIO.write_matrix_binary(filesOut[i], red[1], 'i')"
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}