Plot an Image in 3D with pyplot

Plot Image in 3D

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Oct 24 19:33:19 2019

Plot an image in 3D using pyplot. 
Assumes image is grayscale and in HWC format.

@author: Benedict Wilkins
"""

import matplotlib.pyplot as plt
import numpy as np

img = np.random.uniform(size=(10,10,1))

X = np.arange(0, img.shape[1])
Y = np.arange(0, img.shape[0])
X, Y = np.meshgrid(X, Y)
Z = img.squeeze()

fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(X, Y, Z)