diff --git a/platipy/imaging/generation/augment.py b/platipy/imaging/generation/augment.py
index 6a2a8b22..d8fca7ec 100644
--- a/platipy/imaging/generation/augment.py
+++ b/platipy/imaging/generation/augment.py
@@ -16,6 +16,7 @@
 from collections.abc import Iterable
 import random
 import logging
+import sys
 
 from pathlib import Path