Add data with better representations of the classes and update main.py accordingly.