The model was trained on the
Pokemon TCG — All Image Cards dataset from Kaggle, which contains thousands of official Pokémon card images spanning multiple generations and sets.
1import torch
2import torch.nn as nn
3from torchvision import transforms
4from PIL import Image
5
6# Define your U-Net architecture (must match training)
7# model = UNet(...)
8
9# Load weights
10model.load_state_dict(torch.load("pokemon_unet_colorizer.pth", map_location="cpu"))
11model.eval()
1from PIL import Image
2import torchvision.transforms.functional as TF
3
4# Load and preprocess a grayscale image
5img = Image.open("your_card.png").convert("L") # Grayscale
6img_tensor = TF.to_tensor(img).unsqueeze(0) # Shape: [1, 1, H, W]
7
8# Predict
9with torch.no_grad():
10 output = model(img_tensor) # Shape: [1, 3, H, W]
11
12# Save result
13result = TF.to_pil_image(output.squeeze(0).clamp(0, 1))
14result.save("colorized_card.png")
This model is released under the
Apache 2.0 license.
The training dataset is subject to its own
Kaggle license.
Pokémon and all related names are trademarks of Nintendo / Game Freak / The Pokémon Company. This project is not affiliated with or endorsed by them.
