Updated code for the UI

Some changes were lost when the UI elements were add. Those changes have
been added.