diff --git a/Logic/TrashRecognition/AugmentImages.py b/Logic/TrashRecognition/AugmentImages.py
index 931ab2d..dd26068 100644
--- a/Logic/TrashRecognition/AugmentImages.py
+++ b/Logic/TrashRecognition/AugmentImages.py
@@ -1,16 +1,21 @@
-import Augmentor
+def main():
+     import Augmentor
+     p = [Augmentor.Pipeline("Images/TrainingImages/glass"),
+          Augmentor.Pipeline("Images/TrainingImages/metal"),
+          Augmentor.Pipeline("Images/TrainingImages/paper"),
+          Augmentor.Pipeline("Images/TrainingImages/plastic")]
+     for i in range(len(p)):
+          # Setting rotate values
+          p[i].rotate(probability=0.7, max_left_rotation=10, max_right_rotation=10)
 
-p = [Augmentor.Pipeline("Images/TrainingImages/glass"),
-     Augmentor.Pipeline("Images/TrainingImages/metal"),
-     Augmentor.Pipeline("Images/TrainingImages/paper"),
-     Augmentor.Pipeline("Images/TrainingImages/plastic")]
+          # Setting zoom values
+          p[i].zoom(probability=0.5, min_factor=1.1, max_factor=1.5)
+     
+          # Size of training sample
+          p[i].sample(1000)
 
+          p[i].process()
 
-for i in range(len(p)):
-    p[i].rotate(probability=0.7, max_left_rotation=10, max_right_rotation=10)
-    p[i].zoom(probability=0.5, min_factor=1.1, max_factor=1.5)
-
-    p[i].sample(1000)
-
-    p[i].process()
+if __name__ == "__main__":
+    main()
 
diff --git a/readme.md b/readme.md
index bb593c8..f7fda9c 100644
--- a/readme.md
+++ b/readme.md
@@ -11,7 +11,7 @@
 - [NumPy](http://www.numpy.org/)
 - [PyQt5](https://pypi.org/project/PyQt5/)
 - [TensorFlow](https://www.tensorflow.org/)
-- [Augmentor](https://github.com/mdbloice/Augmentor)
+- [Augmentor](https://augmentor.readthedocs.io)
 
 **Credits**
 - As a learning set of images I used images from [this](https://github.com/garythung/trashnet) project.
\ No newline at end of file