Matplotlib 3D imshow

imshow-3d.py

# Copyright 2023 Nicolas P. Rougier - BSD 2 Clauses licence
# This is an example of 3D projected images with matplotlib using imshow

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cbook as cbook
from matplotlib.path import Path
from mpl_toolkits.mplot3d import proj3d
from matplotlib.patches import PathPatch
import matplotlib.transforms as mtransforms

def warp(T1, T2):
    """
    Return an affine transform that warp triangle T1 into triangle T2.

    Raises
    ------

    `LinAlgError` if T1 or T2 are degenerated triangles
    """
    
    T1 = np.c_[np.array(T1), np.ones(3)]
    T2 = np.c_[np.array(T2), np.ones(3)]
    M = np.linalg.inv(T1) @ T2
    return mtransforms.Affine2D(M.T)

def textured_triangle(ax, T, UV, texture, interpolation="none", image=None):
    """
    Draw a textured triangle T using UV coordinates and given texture.
    
    Parameters
    ----------
    T : (3,2) np.ndarray
      Positions of the triangle vertices
    UV : (3,2) np.ndarray
      UV coordinates of the triangle vertices
    texture: 
      Image to use for texture
    image: AxesImage
      Image if the triangle has been already drawn on axis
    """

    w,h = texture.shape[:2]
    Z = UV*(w,h)
    xmin, xmax = int(np.floor(Z[:,0].min())), int(np.ceil(Z[:,0].max()))
    ymin, ymax = int(np.floor(Z[:,1].min())), int(np.ceil(Z[:,1].max()))
    texture = (texture[ymin:ymax, xmin:xmax,:]).astype(np.uint8)
    extent = xmin/w, xmax/w, ymin/h, ymax/h
    transform = warp (UV,T) + ax.transData
    path =  Path([UV[0], UV[1], UV[2], UV[0]], closed=True)

    if image is not None:
        image.set_transform(transform)
        image.set_extent(extent)
        image.set_clip_path((path,transform))
    else:
        image = ax.imshow(texture, interpolation=interpolation, origin='lower',
                          extent=extent, transform=transform, clip_path=(path,transform))
    return image


with cbook.get_sample_data('grace_hopper.jpg') as image_file:
    texture = plt.imread(image_file)
    
fig = plt.figure(figsize=(8,8))
ax = fig.add_axes([0,0,1,1], projection="3d")


triangles = np.zeros((6,3), dtype = [("xyz", float, 3),
                                     ("uv",  float, 2)])

# XY plane
triangles[0] = [ ((0,0,0), (0,1)),
                 ((1,0,0), (1,1)),
                 ((1,1,0), (1,0)) ]
triangles[1] = [ ((0,0,0), (0,1)),
                 ((1,1,0), (1,0)),
                 ((0,1,0), (0,0)) ]

# XZ plane
triangles[2] = [ ((0,1,0), (0,1)),
                 ((1,1,0), (1,1)),
                 ((1,1,1), (1,0)) ]
triangles[3] = [ ((0,1,0), (0,1)),
                 ((1,1,1), (1,0)),
                 ((0,1,1), (0,0)) ]

# YZ plane
triangles[4] = [ ((0,0,0), (0,1)),
                 ((0,1,0), (1,1)),
                 ((0,1,1), (1,0)) ]
triangles[5] = [ ((0,0,0), (0,1)),
                 ((0,1,1), (1,0)),
                 ((0,0,1), (0,0)) ]

images = [None]*len(triangles)

# For some unknown reason, we need to have a first useless imshow
ax.imshow([[0]], extent=[-1.00,-1.001,-1.00,-1.001])

# Set out axis limit (in relation with our triangle coordinates)
ax.set_xlim(0,1), ax.set_ylim(0,1), ax.set_zlim(0,1)

def update(event=None):
    """
    Update all triangles
    """
    
    for i, (triangle, image) in enumerate(zip(triangles, images)):
        uv = triangle["uv"]
        xy = []
        for xyz in triangle["xyz"]:
            x,y,z = np.array(proj3d.proj_transform(*xyz, ax.get_proj()))
            xy.append((x,y))
        xy = np.array(xy)        
        images[i] = textured_triangle(ax, xy, uv, texture, interpolation="none", image=images[i])

def redraw(event=None):
    update()
    fig.canvas.draw()

update()    
fig.canvas.mpl_connect('draw_event', update)
fig.canvas.mpl_connect('button_release_event', redraw)
# plt.savefig("imshow-3d.png")
plt.show()