UI fix
Fixes for #628