Wrong color information when plotting stylegan3 generated images with motplotlib

[clbr]

Hey there!
I’ve been playing around with plotting images generated from StyleGAN3, but there seems to be something funky going on with the color channels that i can’t really pin down.

Generating an image with:

device = torch.device('cuda')
# open pickle file and load generator
network_pkl = './models/stylegan3-r-afhqv2-512x512.pkl'
with dnnlib.util.open_url(network_pkl) as f:
    G = legacy.load_network_pkl(f, force_fp16=True)['G_ema'].cuda()

noise_mode = 'const'
trunc = 1
seed = 10
z = torch.from_numpy(np.random.RandomState(seed).randn(1, G.z_dim)).to(device)
# returns [b, c, 512, 512] tensor
img = G(z, truncation_psi=trunc, noise_mode=noise_mode, c=None)
# permute to [b, 512, 512, c] and scale to 0-255
img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
# remove batch channel
img = img.squeeze(0) 
# move tensors to cpu
img_cpu = img.cpu()
# plot with pyplot
plt.imshow(img_cpu.numpy())
# save with pillow
PIL.Image.fromarray(img_cpu.numpy(), 'RGB').save(f'seed{seed}.png')

From my understanding the generator of StyleGAN3 returns an image as a tensor of [batch_size, channels, height, width] with channels being RGB. The plotted image seems to be have completely wrong color information.

From plt.imshow(img_cpu.numpy()):

Saving with PIL.Image.fromarray(img_cpu.numpy(), 'RGB').save(f'seed{seed}.png'):

I swapped around the channels to BGR, GRB etc. but couldn’t get the plot to display the correct image. Would appreciate any help with this :).

Cheers

Edit: Issue wasn’t code related. Was using jetbrains DataSpell ide which is in early preview and still has some bugs.

[ptrblck]

That’s quite strange, as I can load your second image via PIL.Image.open and visualize it in PIL as well as matplotlib to get the same output.
Did you change any colormaps in matplotlib or loaded any other template (e.g. via seaborn)?

[clbr]

Nope, haven’t changed anything.

import PIL.Image as Image
import matplotlib.pyplot as plt

with Image.open("seed10.png") as im:
    plt.imshow(im)

This gives me the same image. With seed10.png being the correct image above.

Also tried grabbing the screencapture above as .jpeg, gives me the same result.
This is with matplotlib 3.4.2 and pillow 8.4.0.

But i guess this isn’t really a pytorch issue anymore.

[ptrblck]

Yeah, it’s unrelated to PyTorch, but still interesting as I haven’t seen the issue before.
It doesn’t look like a perfectly inverted image though.
Using your code snippet I get the same image back and inverting it naively via:

plt.imshow(np.array(img)*(-1)+255)

gives

[ptrblck]

Could you run this test and check the colors in your setup, please?

red = np.concatenate((
    np.full((224, 224, 1), 255).astype(np.uint8),
    np.zeros((224, 224, 1)).astype(np.uint8),
    np.zeros((224, 224, 1)).astype(np.uint8)
), axis=2)

green = np.concatenate((
    np.zeros((224, 224, 1)).astype(np.uint8),
    np.full((224, 224, 1), 255).astype(np.uint8),
    np.zeros((224, 224, 1)).astype(np.uint8)
), axis=2)

blue = np.concatenate((
    np.zeros((224, 224, 1)).astype(np.uint8),
    np.zeros((224, 224, 1)).astype(np.uint8),
    np.full((224, 224, 1), 255).astype(np.uint8),
), axis=2)

plt.imshow(red)
plt.imshow(green)
plt.imshow(blue)

[clbr]

Looks like i found the culprit . Was using jetbrains early preview of DataSpell.

DataSpell:

PyCharm:

[ptrblck]

Good to hear you’ve found the issue.
I’m not familiar with this IDE, but do you have an idea why colors seems to be inverted or at least remapped? That seems to be unexpected (unless it’s of course a bug, then it’ll get fixed I guess).

EDIT: Also, just out of interest: could you still post the red, green, blue outputs?

[clbr]

If i would have to guess its probably a bug with how the code editor is rendered.
Cheers for you effort btw. i’ll post the outputs in a sec

Edit: here the the individual outputs in DataSpell, excuse the rather large image.

Seems like they don’t have issues rendering the individual colors, no clue then.

[ptrblck]

Yeah, weird indeed. I was thinking if the color scheme might interact with matplotlib and override the colors or rather be globally applied to the entire IDE. In any case, you might forward the issue to the IDE devs.

[clbr]

Will do. And actually you are right. With the light theme the image is plotted correctly.

[ptrblck]

OK, cool.
To your previous pictures: note that the colors are also not “true” RGB. Look at this quick and dirty check:

[clbr]

Well…
Seems to be correct here as well:

[WoodieDudy]

It also can be fixed by unchecking “invert image outputs for the dark themes” box

[barista (Sascha)]

Oh man, this feature which is active by default has cost me hours of my life. Why would I want the image outputs to be inverted? It’s great to have more features but they should leave this off by default.

[ptrblck]

Yes, I totally agree and don’t know the reason why image outputs should be remapped using the IDE color scheme by default (it also took some time and imagination to figure out what might be causing the issue).
Maybe @clbr forwarded the feedback to the IDE devs so that the default might be changed?