Improved colorization. Still buggy though.