|
--- |
|
license: mit |
|
language: |
|
- en |
|
- fr |
|
tags: |
|
- computer vision |
|
- deep learning |
|
- pytorch |
|
--- |
|
# README for Image Colorization Model |
|
|
|
Welcome to the GitHub repository for our Image Colorization model! This model uses deep learning to add vibrant colors to grayscale images, particularly focusing on historical photographs. Below you'll find instructions on how to set up and use this model. |
|
|
|
## Requirements |
|
To use this model, you need the following: |
|
- Python 3.8 or later |
|
- PyTorch |
|
- torchvision |
|
- PIL (Pillow) |
|
- scikit-image |
|
- matplotlib |
|
- numpy |
|
|
|
You can install the necessary libraries using `pip`: |
|
```bash |
|
pip install torch torchvision Pillow scikit-image matplotlib numpy |
|
``` |
|
|
|
## Model Architecture |
|
`ColorizationNet` is a neural network model built using PyTorch. It combines the power of ResNet and an upsampling network to colorize grayscale images. |
|
|
|
## Setup |
|
First, clone this repository and navigate to the directory: |
|
```bash |
|
git clone [repo-link] |
|
cd [repo-directory] |
|
``` |
|
|
|
Download the trained model weights (`colorization_md1.pth`) and place them in the repository's root directory. |
|
|
|
## Usage |
|
|
|
To colorize an image, follow these steps: |
|
|
|
1. **Prepare the Image**: Ensure your image is in grayscale format. You can use any standard image format like JPG or PNG. |
|
|
|
2. **Run the Model**: Use the following code to colorize your image. |
|
|
|
```python |
|
from colorize import ColorizationNet, colorize_single_image |
|
import torch |
|
import torch.nn as nn |
|
|
|
# Load the model |
|
model = ColorizationNet() |
|
model_path = 'colorization_md1.pth' |
|
pretrained = torch.load(model_path, map_location=lambda storage, loc: storage) |
|
model.load_state_dict(pretrained) |
|
model.eval() # Set the model to evaluation mode |
|
|
|
# Set the path to your image |
|
image_path = 'path_to_your_image.jpg' |
|
save_dir = 'dir_to_save_colorized_image' |
|
criterion = nn.MSELoss() |
|
|
|
# Colorize the image |
|
use_gpu = torch.cuda.is_available() |
|
colorize_single_image(image_path, model, criterion, save_dir, epoch=0, use_gpu=use_gpu) |
|
``` |
|
|
|
3. **View the Results**: The colorized image will be saved in the specified `save_path`. You will find both the original grayscale and the colorized version there. |
|
|
|
## Contributions and Issues |
|
Feel free to contribute to this project by submitting pull requests. If you encounter any issues or have questions, please submit them in the issues section of this repository. |
|
|
|
## License |
|
|
|
|
|
|
|
Enjoy exploring the vibrant colors of history with our model! 🎨📸 |
