Fix small ui issues