update classes in main.py