parvpareek/polygon-colorizer-unet

Polygon Colorizer UNet

This is a CrossAttnUNet trained from scratch to color polygon outlines based on a color name.

How to Use

This model takes a 3-channel polygon outline and a color ID (0-7) to generate a colored polygon. You must have the custom model classes (PolygonColorizer, CrossAttnUNet, etc.) defined in a model.py file to load it.

from model import PolygonColorizer
from huggingface_hub import hf_hub_download
import torch

# Download the weights
weights_path = hf_hub_download(repo_id="parvpareek/polygon-colorizer-unet", filename="pytorch_model.bin")

# Define the model's configuration
config = {'in_ch': 3, 'out_ch': 3, 'base_ch': 64, 'color_embed_dim': 32}
num_colors = 8
embed_dim = 32

# Instantiate the model
model = PolygonColorizer(unet_config=config, num_colors=num_colors, embed_dim=embed_dim)
model.load_state_dict(torch.load(weights_path))

# Now you can use the model for inference!