Fix some import errors after reorganizing