Refactoring distortion correction code