Refactor of flow control for small image edge cases