Cleaning up the structure of the image classification module.