Describe the bug

I have a script where I log an image using wandb.Image then optionally use matplotlib to plot points on top of the same image. Since point plotting is optional, I import matplotlib inside my point plotting method. In practice, this looks like something like this:

def plot_points(img, pts, step):
  import matplotlib.pyplot as plt
  fig, ax = plt.subplots()
  plt.imshow(img)
  plt.scatter(pts[..., 0], pts[..., 1], marker='.')
  wandb.log({'Plot': plt}, step=step)

wimg = wandb.Image(img, caption='Image')
wandb.log({'Image': wimg}, step=step)
if log_pts:
  plot_points(img, pts, step)

When I run this code on my laptop, this works just fine. However, when I run it on my university's computing cluster, I get AttributeError: module 'PIL' has no attribute 'Image' at the plt.imshow call.

I have no idea why this doesn't fail on all of my machines, but I've tracked the issue down to the import order for matplotlib.pyplot. I have no issues if I import pyplot before using wandb.Image, but I get the AttributeError if I import it afterwards. The following is a minimal code example which reproduces the issue:

import numpy as np
import wandb

# import before calling wandb.Image and all is fine 
# import matplotlib.pyplot as plt

wandb.init()

step=0
img = np.zeros((768, 768))

wimg = wandb.Image(img)
wandb.log({'Plot': wimg}, step=step)

import matplotlib.pyplot as plt
fig, ax = plt.subplots()
plt.imshow(img, cmap='gray')
wandb.log({'Image Plot': plt}, step=step)

# paths redacted
Traceback (most recent call last):
  File "/home/.../minimal.py", line 18, in <module>
    plt.imshow(img, cmap='gray')
  File "/usr/local/.../.venv/lib/python3.10/site-packages/matplotlib/pyplot.py", line 3346, in imshow
    __ret = gca().imshow(
  File "/usr/local/.../.venv/lib/python3.10/site-packages/matplotlib/__init__.py", line 1465, in inner
    return func(ax, *map(sanitize_sequence, args), **kwargs)
  File "/usr/local/.../.venv/lib/python3.10/site-packages/matplotlib/axes/_axes.py", line 5751, in imshow
    im.set_data(X)
  File "/usr/local/.../.venv/lib/python3.10/site-packages/matplotlib/image.py", line 721, in set_data
    if isinstance(A, PIL.Image.Image):
AttributeError: module 'PIL' has no attribute 'Image'

The obvious workaround is to just import pyplot at the start of my script, and I'll probably do that now, but I ultimately want to avoid this since point plotting is optional functionality.

Additional Files

No response

Environment

WandB version: 0.15.10

OS: Ubuntu 22.04 (actually the Docker image nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04)

Python version: 3.10.12

Versions of relevant libraries:

• matplotlib 3.8.0
• numpy 1.25.2
• Pillow 10.0.0

I do not have the issue on my laptop with the following configuration:
WandB version: 0.15.10
OS: macOS Ventura 13.5.1 (22G90)
Python version: 3.10.13 (Homebrew)
matplotlib 3.8.0
Pillow 10.0.0

Additional Context

No response