VGreatVig07/color_polygon_unet

🎨 Color Polygon UNet Model

This model is a conditional U-Net trained to generate filled polygon images based on both shape and color. It takes two inputs:

• A binary polygon mask image (e.g., a triangle, star, etc.)
• A color condition (e.g., red, blue, yellow)

It outputs a colorized version of the shape according to the given color prompt.

🧠 Model Architecture

• Backbone: U-Net with encoder-decoder structure
• Input channels: 6 (3 for shape mask, 3 for color hint)
• Output: RGB image of the filled shape

📁 Files

• color_polygon_unet.pth: Trained PyTorch model weights

🔧 How to Use

from huggingface_hub import hf_hub_download
import torch

# Download the model file
model_path = hf_hub_download(
    repo_id="your-username/color-polygon-unet-model",
    filename="color_polygon_unet.pth"
)

# Load the model (assuming you have the UNet class defined)
model = UNet(in_channels=6, out_channels=3)
model.load_state_dict(torch.load(model_path, map_location='cpu'))
model.eval()

🧪 Intended Use

This model is designed for:

• Educational purposes
• Learning conditional image generation
• Demonstrating shape and color-controlled generative models

🛑 Limitations

• Trained on synthetic polygon data
• Not suitable for real-world generalization
• Only supports limited shapes and colors as trained (e.g., red star, blue circle)

📜 License

MIT License

👤 Author

Vighnesh M S