This is a conditional UNet model trained to generate colored polygon images. Given a grayscale/outline polygon image and a color specification, the model generates a colored version of the polygon.

Supported Colors

Blue, Cyan, Green, Magenta, Orange, Purple, Red, Yellow

Model Architecture

Base: Conditional UNet with color embedding
Input Size: 128x128 pixels
Color Embedding: 32 dimensions
Features: 64 base features

Training Details

Framework: PyTorch
Loss Function: Combined MSE + L1 Loss
Optimizer: Adam (lr=0.0001)
Batch Size: 8
Epochs: 50

Usage

python

import torch
from huggingface_hub import hf_hub_download

# Download model
model_path = hf_hub_download(
    repo_id="PentesterPriyanshu/colored-polygon-unet",
    filename="best_model.pth"
)

# Load checkpoint
checkpoint = torch.load(model_path, map_location='cpu')

# Initialize model (you'll need the model architecture code)
model = ConditionalUNet(
    in_channels=3,
    out_channels=3,
    features=64,
    num_colors=8,
    color_embed_dim=32
)
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

Performance

Trained on Kaggle with GPU acceleration
Typical PSNR: 25-35 dB
Fast inference: ~0.1s per image

Limitations

Fixed 128x128 input size
Limited to 8 predefined colors
Works best with simple polygon shapes