Stylegan2-ada tfrecords – ValueError: axes don’t match array, images will work one run and not work the next

I’m working on training a GAN through Google Colab with a dataset of photos I scraped from Wikiart and converted to 1024×1024, but keep getting this error when creating the tfrecords:

Traceback (most recent call last):
  File "dataset_tool.py", line 1249, in <module>
    execute_cmdline(sys.argv)
  File "dataset_tool.py", line 1244, in execute_cmdline
    func(**vars(args))
  File "dataset_tool.py", line 714, in create_from_images
    img = img.transpose([2, 0, 1]) # HWC => CHW
ValueError: axes don't match array

I set it to print out what files it would stall on, and began taking those out of the data set; but what it stalls on seems completely random. It’ll iterate over a file perfectly fine on one run, then fail on it next run after some other troublesome photo is taken out of the dataset.

I’m not sure if the process of constantly removing photos that stall it would ever end/leave me with a meaningful dataset, is there something I should try to fix it?

Answer

Figured out the solution for this, turns out some of the images I scraped were grayscale. To solve this I used imagemagick (also used to resize the photos to 1024×1024) to check the colorspace. I pointed the terminal to the image folder and ran:

magick identify *.jpg

From here, I ctrl+f’ed to see which ones were marked as “Gray” instead of “sRGB”. After taking those out of the dataset it worked like a charm.