Data about images' dimensions can be added. Also add a check if images exist in the dataset, so that changes in the dataset won't be breaking.