refactor color jitter

let's see if we can merge these