---
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

```python
pip install -r requirements.txt
```

### 2️⃣ Load the Model
```python
import torch
from model import ImprovedUNet
```

# Create the model instance
```python
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
```python
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:
![Colorized Example](test_result_PARIS.png)


✨ Author
This model was developed by Eric Houzelle.