Modified draw code