GAN Based colouring model
Project description
ColourGAN
Simple PyTorch based GAN for colouring black and white images.
Installation
The code is available as python package.
Run following command for installing colourgan with all it's dependencies.
pip install colourgan
Additionally for trying out streamlit interface, please install streamlit using following command
pip install streamlit
You can also install it from this repository.
git clone https://github.com/narenderkumarnain/ColourGAN.git
cd ColourGAN
python setup.py install
Usage
Inference
For trying our pretrained model, please download the
Generator Weight File
Now use these lines of code for inference on sample image.
#imports
import cv2
from colourgan.model import ColourGAN
from colourgan.config import get_cfg
# preparing the config
cfg = get_cfg()
# adding path to generator weights
cfg.initial_weights_generator = 'checkpoint_ep99_gen.pt'
# loading the model
model = ColourGAN(cfg , inference=True)
# loading the image
img = cv2.imread('sample.png')
# running inference
res = model.inference(img)
# saving result image
cv2.imwrite('sample_res.png' , res)
Streamlit interface
For trying out the streamlit interface, create a file app.py with following code.
# import
from colourgan.streamlit import app
# path to weights file
path_to_generator_weights = 'checkpoint_ep99_gen.pt'
if __name__ == '__main__':
app(path_to_generator_weights)
We expect you have already installed streamlit.
Use followind code to run the application
streamlit run app.py
Training
Use following Lines of code for training the model on Cifar10 dataset.
# imports
import torch
from colourgan.config import get_cfg
from colourgan.data import Cifar10Dataset
from colourgan.model import ColourGAN
# Downloading the data and creating the dataloaders
dataloaders = Cifar10Dataset('cifar10', batch_size = 8).get_dataloaders()
# config
cfg = get_cfg()
# creating the model
model = ColourGAN(cfg)
# train
model.train(dataloaders['train'],dataloaders['test'], epochs = 100)
Pretrained Weights
References
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for colourgan-1.0.3-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 6ac07e671acd178f9fff1a0613187b04a934f8be71ee4918faa31a83d37e5993 |
|
MD5 | 5bedcee4283efafe4faf5abc6a7579ea |
|
BLAKE2b-256 | 486c6c7358f7cdc42ee4112dafd5dac9b354ed3faf39abd99b48d381a1ceef41 |