ADD: Fix wrong class names