Add further UI improvements