diff --git a/README.md b/README.md
index 52bc03e..4741c5d 100644
--- a/README.md
+++ b/README.md
@@ -53,6 +53,8 @@ There are a few supported operations:
 * `skew` - random image skewing followed by `-p` option and constant `magnitude` value of `0.7`
 * `zoom` - random image zooming followed by `-p`,`-minf` and `-maxf` options
 
+Output images will be saved in `output` dir of the root.
+
 # Donate
 Just put a star on this repository 🌞 Thanks!
 
diff --git a/bin/rudi b/bin/rudi
index 8968ea1..74bdcb2 100644
--- a/bin/rudi
+++ b/bin/rudi
@@ -1,6 +1,4 @@
 #!/usr/bin/env python
 
 import rudi
-
-if __name__ == "__main__":
-    rudi.cli()
\ No newline at end of file
+rudi.cli()
\ No newline at end of file
diff --git a/rudi/cli.py b/rudi/cli.py
index 3452056..c701876 100644
--- a/rudi/cli.py
+++ b/rudi/cli.py
@@ -1,4 +1,5 @@
 from .utils import convert_
+from .utils import augment_
 import click
 
 
@@ -38,8 +39,11 @@ def convert(type, size, root_dir):
 @click.option('--zoom/--no-zoom', default=True)
 
 @click.argument('root_dir')
-def augment(prob, num, flip, root_dir):
+def augment(**args):
     """
     Dataset augmentation
     """
+    augment_(args)
     pass
+
+
diff --git a/rudi/utils/__init__.py b/rudi/utils/__init__.py
index 12cf5ce..4d93375 100644
--- a/rudi/utils/__init__.py
+++ b/rudi/utils/__init__.py
@@ -1 +1,2 @@
-from .converter import *
\ No newline at end of file
+from .converter import *
+from .augmentor import *
\ No newline at end of file
diff --git a/rudi/utils/augmentor.py b/rudi/utils/augmentor.py
new file mode 100644
index 0000000..2931fc1
--- /dev/null
+++ b/rudi/utils/augmentor.py
@@ -0,0 +1,35 @@
+import Augmentor
+
+def augment_(args):
+
+    p = Augmentor.Pipeline(args["root_dir"])
+
+    #Default params
+    probability_ = args["probability"]
+    magnitude_ = args["magnitude"]
+    left_r_ = args["left_r"]
+    right_r = args["right_r"]
+    grid_ = args["grid"]
+    minf_ = args["min_factor"]
+    maxf_ = args["max_factor"]
+    num_ = args["num"]
+
+    # Operations
+    flip_ = args["flip"]
+    rotate_ = args["rotate"]
+    distortion_ = args["distortion"]
+    skew_ = args["skew"]
+    zoom_ = args["zoom"]
+
+    if rotate_:
+        p.rotate(probability=probability_, max_left_rotation=left_r_, max_right_rotation=right_r)
+    if distortion_:
+        p.random_distortion(probability=probability_, grid_width=grid_, grid_height=grid_, magnitude=magnitude_)
+    if skew_:
+        p.skew(probability=probability_, magnitude=0.7)
+    if zoom_:
+        p.zoom(probability=probability_, min_factor=minf_, max_factor=maxf_)
+    if flip_:
+        p.flip_random(probability=probability_)
+
+    p.sample(num_)
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 7217b11..33bee17 100644
--- a/setup.py
+++ b/setup.py
@@ -6,7 +6,7 @@
 setuptools.setup(
     name='rudi',
     version='1.0',
-    scripts=['bin/rudi'],
+    # scripts=['bin/rudi'],
     author="Petro Liashchynskyi",
     description="Small, fast and simple Python CLI image converter for CNNs.",
     long_description=long_description,
@@ -15,7 +15,7 @@
     packages=setuptools.find_packages(),
     install_requires=[
         'click',
-        #'Augmentor',
+        'Augmentor',
         'Pillow'
     ],
     classifiers=[
@@ -23,4 +23,7 @@
         "License :: OSI Approved :: MIT License",
         "Operating System :: OS Independent",
     ],
+    entry_points={
+        'console_scripts': ['rudi=rudi:cli'],
+    }
 )