Edited UI a little