Improvements in data augmentation

rvmaretto · May 25, 2018 · fb675df · fb675df
1 parent 15f13b3
commit fb675df
Show file tree

Hide file tree

Showing 5 changed files with 284 additions and 115 deletions.
diff --git a/notebooks/generate_samples_one_scene.ipynb b/notebooks/generate_samples_one_scene.ipynb
diff --git a/notebooks/plot_preprocessor_results.ipynb b/notebooks/plot_preprocessor_results.ipynb
@@ -45,9 +45,7 @@
   {
    "cell_type": "code",
    "execution_count": 12,
-   "metadata": {
-    "collapsed": false
-   },
+   "metadata": {},
    "outputs": [
     {
      "name": "stdout",
@@ -82,9 +80,7 @@
   {
    "cell_type": "code",
    "execution_count": 13,
-   "metadata": {
-    "collapsed": false
-   },
+   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -120,9 +116,7 @@
   {
    "cell_type": "code",
    "execution_count": 14,
-   "metadata": {
-    "collapsed": false
-   },
+   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -142,9 +136,7 @@
   {
    "cell_type": "code",
    "execution_count": 15,
-   "metadata": {
-    "collapsed": false
-   },
+   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -181,9 +173,7 @@
   {
    "cell_type": "code",
    "execution_count": 16,
-   "metadata": {
-    "collapsed": false
-   },
+   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -203,9 +193,7 @@
   {
    "cell_type": "code",
    "execution_count": 17,
-   "metadata": {
-    "collapsed": false
-   },
+   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -240,9 +228,7 @@
   {
    "cell_type": "code",
    "execution_count": 18,
-   "metadata": {
-    "collapsed": false
-   },
+   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -262,9 +248,7 @@
   {
    "cell_type": "code",
    "execution_count": 19,
-   "metadata": {
-    "collapsed": false
-   },
+   "metadata": {},
    "outputs": [
     {
      "name": "stdout",
@@ -303,7 +287,7 @@
  "metadata": {
   "anaconda-cloud": {},
   "kernelspec": {
-   "display_name": "Python [default]",
+   "display_name": "Python 3",
    "language": "python",
    "name": "python3"
   },
@@ -317,7 +301,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.5.3"
+   "version": "3.5.4"
   }
  },
  "nbformat": 4,

diff --git a/notebooks/rasterize_and_plot.ipynb b/notebooks/rasterize_and_plot.ipynb
@@ -45,9 +45,7 @@
   {
    "cell_type": "code",
    "execution_count": 3,
-   "metadata": {
-    "collapsed": false
-   },
+   "metadata": {},
    "outputs": [
     {
      "name": "stdout",
@@ -68,9 +66,7 @@
   {
    "cell_type": "code",
    "execution_count": 8,
-   "metadata": {
-    "collapsed": false
-   },
+   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -104,9 +100,7 @@
   {
    "cell_type": "code",
    "execution_count": 25,
-   "metadata": {
-    "collapsed": false
-   },
+   "metadata": {},
    "outputs": [
     {
      "name": "stdout",
@@ -132,9 +126,7 @@
   {
    "cell_type": "code",
    "execution_count": 26,
-   "metadata": {
-    "collapsed": false
-   },
+   "metadata": {},
    "outputs": [
     {
      "name": "stdout",
@@ -185,7 +177,7 @@
  "metadata": {
   "anaconda-cloud": {},
   "kernelspec": {
-   "display_name": "Python [default]",
+   "display_name": "Python 3",
    "language": "python",
    "name": "python3"
   },
@@ -199,7 +191,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.5.3"
+   "version": "3.5.4"
   }
  },
  "nbformat": 4,

diff --git a/src/datasetGen/data_augment.py b/src/datasetGen/data_augment.py
@@ -0,0 +1,56 @@
+import tensorflow as tf
+import math
+import numpy as np
+
+# Methods based on https://medium.com/ymedialabs-innovation/data-augmentation-techniques-in-cnn-using-tensorflow-371ae43d5be9#f8ea
+def rotate_images(images, angles, data_type=np.float32):
+    if(not isinstance(angles, list)):
+        angles = [angles]
+
+    if(not isinstance(images, list)):
+        images = [images]
+
+    tf.reset_default_graph()
+
+    tf_shape = (None,) + images[0].shape
+
+    img = tf.placeholder(data_type, shape = tf_shape)
+    radian = tf.placeholder(tf.float32, shape = (len(images)))
+    tf_img = tf.contrib.image.rotate(img, radian)
+
+    with tf.Session() as sess:
+        rotated_imgs = []
+        sess.run(tf.global_variables_initializer())
+        for angle in angles:
+            radian_angle = angle * math.pi / 180
+            radian_list = [radian_angle] * len(images)
+            rot_img = sess.run(tf_img, feed_dict = {img: images, radian: radian_list})
+            rot_img = np.array(rot_img, dtype=data_type)
+            rotated_imgs.extend(rot_img)
+
+        sess.close()
+        return rotated_imgs
+
+def flip_images(images, data_type=np.float32):
+    if(not isinstance(images, list)):
+        angles = [images]
+
+    flipped_imgs = []
+    tf.reset_default_graph()
+
+    tf_img = tf.placeholder(data_type, shape = images[0].shape)
+    tf_img1 = tf.image.flip_left_right(tf_img)
+    tf_img2 = tf.image.flip_up_down(tf_img)
+    tf_img3 = tf.image.transpose_image(tf_img)
+
+    with tf.Session() as sess:
+        sess.run(tf.global_variables_initializer())
+
+        for img in images:
+            flipped = sess.run([tf_img1, tf_img2, tf_img3], feed_dict = {tf_img: img})
+            flipped_imgs.extend(flipped)
+
+        sess.close()
+
+    flipped_imgs = np.array(flipped_imgs, dtype = data_type)
+    return flipped_imgs
diff --git a/src/datasetGen/rasterizer.py b/src/datasetGen/rasterizer.py
@@ -76,7 +76,7 @@ def rasterize_layer(self):
         vector_ds = ogr.Open(self.vector_path)
         vector_layer = vector_ds.GetLayer()
         self.labeled_raster = np.ma.masked_all((self.base_raster.RasterYSize, 
-                                self.base_raster.RasterXSize),
+                                self.base_raster.RasterXSize, 1),
                                 dtype=np.int16)
 
         for lid, label in enumerate(self.class_names):