Skip to content

Instantly share code, notes, and snippets.

@guidorice
Last active August 6, 2023 00:24
Show Gist options
  • Save guidorice/8edbfe1ecb6041af427aef7176ef4c11 to your computer and use it in GitHub Desktop.
Save guidorice/8edbfe1ecb6041af427aef7176ef4c11 to your computer and use it in GitHub Desktop.
Mojo Tensor + plotting
from DType import DType
from Tensor import Tensor, TensorShape
from List import VariadicList
from Random import rand

alias height = 256
alias width = 256
alias channels = 3

# Create the tensor of dimensions height, width, channels
let image = Tensor[DType.float32](TensorShape(height, width, channels))

# Fill the tensor with random values.
rand(image.data(), image.num_elements())

# Declare the grayscale image.
var gray_scale_image = Tensor[DType.float32](TensorShape(height, width))

# Perform the RGB to grayscale transform.
for y in range(height):
    for x in range(width):
        let r = image[y,x,0]
        let g = image[y,x,1]
        let b = image[y,x,2]
        gray_scale_image[VariadicList(y, x)] = 0.299 * r + 0.587 * g + 0.114 * b
def tensor_to_numpy(t: Tensor[DType.float32]) -> PythonObject:
    let np = Python.import_module("numpy")
    let numpy_array = np.zeros((height, width), np.float32)
    for col in range(width):
        for row in range(height):
            numpy_array.itemset((row, col), t[row, col, 0])  # fetch 0th channel
    return numpy_array
def make_plot(t: Tensor[DType.float32]):
    np = Python.import_module("numpy")
    plt = Python.import_module("matplotlib.pyplot")
    colors = Python.import_module("matplotlib.colors")
    dpi = 64
    fig = plt.figure(1, [width, height], dpi)
    ax = fig.add_axes([0.0, 0.0, 1.0, 1.0], False, 1)
    light = colors.LightSource(315, 10, 0, 1, 1, 0)
    arr = tensor_to_numpy(t)
    image = light.shade(arr, plt.cm.hot, colors.PowerNorm(0.3), "hsv", 0, 0, 1.5)
    plt.imshow(image)
    plt.axis("off")
    plt.show()

make_plot(image)
error: execution stopped with unexpected state.
error: Execution was interrupted, reason: shared-library-event.
The process has been left at the point where it was interrupted, use "thread return -x" to return to the state before expression evaluation.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment