Fix selection code to work after repaint