Small refactoring, found problem for images of different dimensions