Added import os statement to main.py