| # /// 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") | |