Refactor image functionality