README.md

metadata

license: mit
tags:
  - image-colorization
  - pytorch
model_name: Simple Colorizer
library_name: pytorch

🎨 Simple Colorizer - Image Colorization Model

This repository contains a PyTorch-trained U-Net model that automatically colorizes grayscale images.

---

📂 Repository Contents

• best_colorization_model.pth: Trained model weights
• model.py: The ImprovedUNet architecture definition
• README.md: This file

---

🚀 Usage Example

1️⃣ Install Dependencies

pip install -r requirements.txt

2️⃣ Load the Model

import torch
from model import ImprovedUNet

Create the model instance

model = ImprovedUNet()

# Load the weights
checkpoint = torch.load("best_colorization_model.pth", map_location="cpu")
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()

3️⃣ Colorize an Image

from PIL import Image
import torchvision.transforms as T

img = Image.open("path/to/grayscale_image.jpg").convert("L")
transform = T.Compose([
    T.Resize((256, 256)),
    T.ToTensor(),
    T.Normalize(mean=[0.5], std=[0.5])
])

input_tensor = transform(img).unsqueeze(0) 

with torch.no_grad():
    output = model(input_tensor)

output_image = output.squeeze(0).permute(1, 2, 0).numpy()
output_image = (output_image * 255).clip(0, 255).astype("uint8")

Image.fromarray(output_image).save("colorized_output.png")

ℹ️ Training Information

Architecture: Custom U-Net (ImprovedUNet)

Input Size: 256x256 pixels

Optimizer: Adam

Loss Function: MSE

Epochs: [Specify the number of epochs]

📈 Results Here is an example of an image colorized by the model:

✨ Author This model was developed by Eric Houzelle.