File size: 759 Bytes
63f348a 5ddd783 63f348a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
# /// script
# requires-python = "==3.10"
# dependencies = [
# "kernels",
# "numpy",
# "pillow",
# "torch",
# ]
# ///
import torch
from PIL import Image
import numpy as np
from kernels import get_kernel
# This downloads, caches, and loads the kernel library
# and makes the custom op available in torch.ops
img2gray_lib = get_kernel("drbh/img2gray")
img = Image.open("kernel-builder-logo-color.png").convert("RGB")
img = np.array(img)
img_tensor = torch.from_numpy(img).cuda()
print(img_tensor.shape) # HWC
gray_tensor = img2gray_lib.img2gray(img_tensor).squeeze()
print(gray_tensor.shape) # HW
# save the output image
gray_img = Image.fromarray(gray_tensor.cpu().numpy().astype(np.uint8))
gray_img.save("kernel-builder-logo-gray2.png")
|