update readme

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!
 

bin/rudi

@@ -1,6 +1,4 @@
 #!/usr/bin/env python
 
 import rudi
-
-if __name__ == "__main__":
-    rudi.cli()
+rudi.cli()

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
+
+

rudi/utils/__init__.py

@@ -1 +1,2 @@
-from .converter import *
+from .converter import *
+from .augmentor import *

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_)

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,12 +15,15 @@
     packages=setuptools.find_packages(),
     install_requires=[
         'click',
-        #'Augmentor',
+        'Augmentor',
         'Pillow'
     ],
     classifiers=[
         "Programming Language :: Python :: 3",
         "License :: OSI Approved :: MIT License",
         "Operating System :: OS Independent",
     ],
+    entry_points={
+        'console_scripts': ['rudi=rudi:cli'],
+    }
 )