PyTorch implementation of a UNet model designed to generate images of colored polygons from two inputs: a polygon image (e.g., triangle, square) and a color name (e.g., "blue", "red")
- If colored_polygon.ipynb doesn't open download and open in colab, (its just not displaying in github) - https://drive.google.com/file/d/1rvh-t-HbY-Q1w2tvQbf8IA0eLqPc7Cvh/view?usp=sharing
- Wandb report (sandy-moon-24) - https://api.wandb.ai/links/gaonkarsub-kogo/j5b21amh
- Drive - Models, wandb etc - https://drive.google.com/drive/folders/1xcU-lhZh5j82ckcwvHnTQrNuzClPq7QM?usp=sharing
Overview The UNet architecture is a fully convolutional network with a symmetric encoder-decoder structure, featuring skip connections to preserve spatial information. For this task, I designed an EnhancedConditionedUNet that incorporates color conditioning via Feature-wise Linear Modulation (FiLM) layers (Perez et al., 2018) and self-attention mechanisms at the bottleneck to enhance feature integration. The model takes a grayscale image (1 channel) and a color index, producing an RGB image (3 channels). Below, I describe each component, its intuition, and the mathematical formulation.
- Input Processing Input: The model accepts a grayscale image (H = W = 128) and a color index Color Embedding:The color index is mapped to a dense embedding vector using an embedding layer followed by a linear transformation, layer normalization, and ReLU activation. The embedding layer converts the discrete color index into a continuous representation, allowing the model to learn meaningful features for each color. The subsequent linear layer and normalization enhance the expressiveness and stability of the embedding.
self.color_embedding = nn.Sequential(
nn.Embedding(n_colors, emb_dim),
nn.Linear(emb_dim, emb_dim),
nn.LayerNorm(emb_dim),
nn.ReLU(inplace=True)
)- Encoder (Contracting Path) Initial convolution(DoubleConv)
self.inc = DoubleConv(n_channels, f, residual=True)The encoder progressively downsamples the input image to extract hierarchical features, reducing spatial dimensions while increasing channel depth. The DoubleConv block consists of two 3x3 convolutions, each followed by batch normalization and ReLU activation, with an optional residual connection. The initial convolution extracts low-level features (e.g., edges, corners) from the grayscale polygon image. The residual connection ($ x + f(x) $) mitigates vanishing gradients and improves training stability, especially for deep networks.
Downsampling Blocks (DownBlock):
self.down1 = DownBlock(f, f*2)
self.down2 = DownBlock(f*2, f*4)
self.down3 = DownBlock(f*4, f*8)
self.down4 = DownBlock(f*8, f*16)
self.down5 = DownBlock(f*16, f*16)
self.down6 = DownBlock(f*16, f*32)Each DownBlock consists of a max-pooling layer (2x2, stride 2) followed by a DoubleConv block. FiLM layers condition the features using the color embedding. Max-pooling reduces spatial dimensions (e.g., from 128x128 to 64x64), allowing the model to capture higher-level features (e.g., polygon shapes) while increasing the receptive field. The FiLM layers modulate features based on the color, enabling the model to incorporate color information early in the network.
- Bottleneck The bottleneck is the deepest part of the network, where features are highly compressed (e.g., 4x4 spatial resolution with 2048 channels for base_filters=64). Self-Attention (SelfAttention):
self.attention1 = SelfAttention(f*32)
self.attention2 = SelfAttention(f*32)
self.dropout = nn.Dropout2d(dropout_rate)Two self-attention layers are applied, with dropout in between. Each layer computes query, key, and value projections, followed by a softmax attention mechanism and a residual connection. Self-attention allows the model to capture long-range dependencies across the image, which is crucial for ensuring consistent color application across the polygon. Dropout prevents overfitting by randomly masking features during training.
self.up0 = UpBlock(f*32, f*16, f*16)
self.up1 = UpBlock(f*16, f*16, f*16)
self.up2 = UpBlock(f*16, f*8, f*8)
self.up3 = UpBlock(f*8, f*4, f*4)
self.up4 = UpBlock(f*4, f*2, f*2)
self.up5 = UpBlock(f*2, f, f)Each UpBlock uses a transposed convolution to upsample the input, concatenates with the corresponding encoder features (skip connection), applies a DoubleConv, and uses FiLM for conditioning. Transposed convolutions increase spatial dimensions (e.g., from 4x4 to 8x8), while skip connections provide high-resolution details from the encoder, ensuring accurate polygon boundaries. FiLM layers maintain color consistency throughout the decoding process.
- Decoder (Expanding Path) The decoder upsamples the bottleneck features, integrating skip connections from the encoder to recover spatial details. Upsampling Blocks (UpBlock):
self.up0 = UpBlock(f*32, f*16, f*16)
self.up1 = UpBlock(f*16, f*16, f*16)
self.up2 = UpBlock(f*16, f*8, f*8)
self.up3 = UpBlock(f*8, f*4, f*4)
self.up4 = UpBlock(f*4, f*2, f*2)
self.up5 = UpBlock(f*2, f, f)Each UpBlock uses a transposed convolution to upsample the input, concatenates with the corresponding encoder features (skip connection), applies a DoubleConv, and uses FiLM for conditioning. Transposed convolutions increase spatial dimensions (e.g., from 4x4 to 8x8), while skip connections provide high-resolution details from the encoder, ensuring accurate polygon boundaries. FiLM layers maintain color consistency throughout the decoding process.
self.outc = nn.Sequential(
nn.Conv2d(f, f, kernel_size=3, padding=1),
nn.BatchNorm2d(f),
nn.ReLU(inplace=True),
nn.Conv2d(f, f//2, kernel_size=3, padding=1),
nn.BatchNorm2d(f//2),
nn.ReLU(inplace=True),
nn.Conv2d(f//2, n_classes, kernel_size=1),
nn.Tanh()
)The final convolution refines the features into a 3-channel RGB image, with Tanh ensuring outputs are in
- Loss Function (CombinedLoss)
class CombinedLoss(nn.Module):
def __init__(self, mse_weight=0.8):
super().__init__()
self.mse_weight = mse_weight
self.mse = nn.MSELoss()
def forward(self, outputs, targets):
mse_loss = self.mse(outputs, targets)
color_loss = torch.mean(torch.abs(outputs.mean(dim=(2, 3)) - targets.mean(dim=(2, 3))))
return self.mse_weight * mse_loss + (1 - self.mse_weight) * color_lossThe loss combines Mean Squared Error (MSE) for pixel-wise accuracy and a color consistency loss to ensure the output image's average color matches the target. The mse_weight balances these objectives.
Dataset (PolygonDataset)
class PolygonDataset(Dataset):
def __init__(self, data_dir, transform=None, color_map=None, augment=True):
self.data_dir = Path(data_dir)
self.inputs_dir = self.data_dir / 'inputs'
self.outputs_dir = self.data_dir / 'outputs'
self.transform = transform
self.augment = augment
json_path = self.data_dir / 'data.json'
with open(json_path, 'r') as f:
self.metadata = json.load(f)
if color_map is None:
self.color_map, self.idx_to_color = self._create_color_map()
else:
self.color_map = color_map
self.idx_to_color = {v: k for k, v in color_map.items()}The dataset loads paired grayscale input images, RGB output images, and color labels from a JSON file. It creates a color map to convert color names to indices, ensuring consistency across training and validation sets.
- Images are resized to 128x128 and converted to tensors.
- Input images are grayscale (1 channel), normalized to
$[-1, 1]$ using mean=0.5, std=0.5. - Output images are RGB (3 channels), normalized similarly.
- Augmentations (flips, rotations, affine transforms, noise, blur) are applied to input-output pairs to improve generalization.
class PairedTransforms:
def __init__(self, prob=0.5, degrees=15, translate=(0.08, 0.08), scale=(0.9, 1.1)):
self.prob = prob
self.degrees = degrees
self.translate = translate
self.scale = scale
self.color_jitter = transforms.ColorJitter(brightness=0.1, contrast=0.1)
def __call__(self, input_img, output_img):
# Random noise, blur, flips, rotation, affine transforms
...Paired augmentations ensure that input and output images undergo the same transformations, preserving their correspondence. This enhances robustness to variations in polygon orientation and position.
for epoch in range(NUM_EPOCHS):
train_loss = train_epoch(model, train_loader, optimizer, criterion, DEVICE)
val_loss = validate(model, val_loader, criterion, DEVICE)
scheduler.step()
wandb.log({"epoch": epoch, "train_loss": train_loss, "val_loss": val_loss, "learning_rate": optimizer.param_groups[0]['lr']})
if val_loss < best_val_loss:
best_val_loss = val_loss
save_model(model, epoch, val_loss, f"{CHECKPOINT_DIR}/best_model.pt")The training loop optimizes the model using AdamW with weight decay and a cosine annealing learning rate scheduler to smoothly reduce the learning rate. Checkpoints are saved for the best validation loss.
params_to_test = {
"base_filters": [32, 48, 64],
"learning_rate": [1e-3, 3e-4, 1e-4],
"batch_size": [8, 16],
"dropout_rate": [0.1, 0.2, 0.3],
"mse_weight": [0.7, 0.8, 0.9]
}I tested combinations of hyperparameters to balance model capacity, training stability, and generalization. Smaller base_filters reduce overfitting, while lower learning rates ensure stable convergence. The mse_weight balances pixel accuracy and color consistency.
The training and testing process for the UNet model was conducted to generate colored polygon images from grayscale inputs, conditioned on color labels. The dataset consisted of 56 training samples and 5 validation samples, with 8 unique colors mapped to indices. The model, implemented with PyTorch, utilized a custom UNet architecture with Feature-wise Linear Modulation (FiLM) for color conditioning and self-attention for global feature integration. Hyperparameter tuning was performed over three configurations, testing variations in base_filters, learning_rate, batch_size, dropout_rate, and mse_weight. The training process used a combined loss function (0.8 MSE + 0.2 color loss) and was monitored using Weights & Biases (wandb) for loss tracking and visualization.
Config 1: base_filters=32, learning_rate=0.001, batch_size=16, dropout_rate=0.1, mse_weight=0.8
Train Loss: 0.2174, Accuracy: 66.07%
Val Loss: 0.1225, Accuracy: 80.00%
Config 2: base_filters=32, learning_rate=0.0003, batch_size=8, dropout_rate=0.2, mse_weight=0.7
Train Loss: 0.4471, Accuracy: 44.64%
Val Loss: 0.4406, Accuracy: 80.00%
Config 3: base_filters=48, learning_rate=0.0003, batch_size=8, dropout_rate=0.2, mse_weight=0.8
Train Loss: 0.2014, Accuracy: 76.79%
Val Loss: 0.1900, Accuracy: 80.00%
Best Configuration: Config 1 achieved the lowest validation loss (0.1225) and was selected for final training.
Results: The best configuration (e.g., base_filters=48, learning_rate=3e-4, batch_size=8, dropout_rate=0.2, mse_weight=0.8) achieved the lowest validation loss, indicating optimal performance.metrics = calculate_metrics(model, val_loader, DEVICE, criterion)
print(f"MSE Loss: {metrics['mse_loss']:.6f}")
print(f"Color Accuracy: {metrics['color_accuracy']:.2%}")
print(f"Classification Report:\n{metrics['classification_report']}")MSE loss measures pixel-wise accuracy, while color accuracy (based on dominant color) evaluates color prediction correctness. The confusion matrix visualizes color prediction errors.
The model was trained for 100 epochs using the best hyperparameters. The training utilized a cosine annealing learning rate scheduler, starting at 0.001 and decaying to 1e-6. The training loss decreased steadily from 1.0863 to 0.0136, and the validation loss improved from 1.9181 to 0.0081, with the best validation loss of 0.00786 at epoch 93. The model was saved at each improvement, and sample predictions were visualized every 5 epochs. Final Evaluation
- MSE Loss: 0.00786
- Color Accuracy: 100.00% (based on dominant RGB channel matching)
- Classification Report: Precision, recall, and F1-score were low (macro avg: 0.12, 0.25, 0.17) due to limited validation samples (5) and imbalanced color distribution, with some colors (e.g., blue, green, yellow) having zero precision due to no predictions.
- Training Dynamics: The consistent decrease in both train and validation losses indicates effective learning, with the cosine scheduler aiding convergence. The model generalized well, as evidenced by the low validation loss.
- Accuracy: The 100% color accuracy suggests the model correctly predicts dominant colors, but the classification report highlights challenges with underrepresented colors, likely due to the small validation set.
- Challenges: The small dataset (56 training, 5 validation samples) limits robustness, and the UndefinedMetricWarning indicates some colors were not predicted, possibly due to class imbalance.
- The wandb logs (run: sandy-moon-24) show smooth loss curves and high-quality sample images, confirming the model’s ability to generate accurate colored polygons. The visualizations saved in checkpoints/ demonstrate qualitative improvements over epochs.
The UNet model, enhanced with FiLM conditioning and self-attention, successfully learned to generate colored polygon images, achieving a low validation loss of 0.00786 and 100% color accuracy with the best hyperparameters (base_filters=32, learning_rate=0.001, batch_size=16, dropout_rate=0.1, mse_weight=0.8). Despite the small dataset size, the model demonstrated robust performance, with consistent loss reduction and high-quality outputs. However, the low precision in some classes suggests a need for a larger, more balanced dataset to improve generalization. The use of wandb facilitated effective monitoring, and the results validate the model’s suitability for conditional image generation tasks. Future work could involve augmenting the dataset and exploring additional conditioning mechanisms to enhance performance on underrepresented colors.