diff --git a/traffic_gym.py b/traffic_gym.py
index daf55f21..831e8de3 100644
--- a/traffic_gym.py
+++ b/traffic_gym.py
@@ -9,6 +9,7 @@
 # from skimage import measure, transform
 # from matplotlib.image import imsave
 import PIL
+from PIL import Image
 from custom_graphics import draw_dashed_line, draw_text, draw_rect
 from gym import core, spaces
 import os