Group Project for IDL
- python >= 3.7
- torch >= 1.8.1
- torchvision >= 0.9.1
# Create environment with Python 3.7 to be compatible with the colorization model
conda create -n idl-project python=3.7 -y
conda activate idl-project
# Since there is no suitable pytorch for Python 3.6,
# let's just use CPU version of Pytorch for inference
pip install torch
pip install pyyaml
pip install torchvision
pip install tabulate
pip install einops
pip install scipy
pip install tqdm
pip install matplotlib
pip install rich
pip install tensorboard
pip install chardet
pip install gdown
pip install wandb
pip install onnxruntime
pip install pyqt5
pip install openai==1.3.7
- Step 1: set up and activate the conda environment following instructions above
- Step 2: create
segmentation/models/
directory, enter it, and rungdown https://drive.google.com/u/1/uc?id=1AcgEK5aWMJzpe8tsfauqhragR0nBHyPh
to get the pre-trained model on full ADE20K dataset - Step 3: enter
colorization
, and rungdown https://drive.google.com/u/2/uc?id=1G7nN4PL48TB6aJjCn7F6tJSM_d7pZ3_4
to get a pre-trained model - Step 4: run
python main.py [--example, for using example color mapping for the existing input.png]
and follow the instructions. Results will be saved inresult/
- Full MIT Scene Parsing ADE 20K Benchmark: http://sceneparsing.csail.mit.edu/
- Our adapted/filtered room-specific dataset: Google Drive
- Download:
gdown https://drive.google.com/u/2/uc?id=1-W-A9gDkVitq7lcGg2srm01DaB5nPsL8
- Unzip and place into
semantic-segmentation/data/
- Download:
- Our filtered room-specific dataset, but converted to sketches: Google Drive
- Download:
gdown https://drive.google.com/u/2/uc?id=1p3D5Y6X89SIOhO7ostJM6R6G1xKPLPxI
- Unzip and place into
semantic-segmentation/data/
- Download:
- Backbone model (Mit-B2) for training and finetuning: Google Drive
- Download:
gdown https://drive.google.com/u/2/uc?id=1Ju_8VWh8aG7mKrvfshRCfwNkk1qYIEvl
- Place into
semantic-segmentation/models/
- Download:
- Pretrained model - original which is provided by the authors and trained on the full original ADE-20K dataset: Google Drive
- Download:
gdown https://drive.google.com/u/1/uc?id=1AcgEK5aWMJzpe8tsfauqhragR0nBHyPh
- Place into
semantic-segmentation/models/
- Download:
- Pretrained model - rooms only which is trained using Mit-B2 backbone on our filtered room-specific dataset
- SegRooms_100Epochs_32mIoU:
- Google Drive
gdown https://drive.google.com/u/2/uc?id=1IQNBteYCB1R7kYUlO-uNDOFh2PR80KCP
- SegRooms_250Epochs_29mIoU:
- Google Drive
gdown https://drive.google.com/u/2/uc?id=1rrt5m3GeWLyvVhkQ2ibXzqQhs5kOQAJT
- SegRooms_100Epochs_32mIoU:
- Pretrained model - room sketches only
- For Training, run
python ./train.py --name [experiment_name]
- Use your own wandb API key in
train.py -> wandb.login()
- Make sure you have CUDA GPU available for Pytorch, otherwise it is extremely slow on CPU
- Recommended to use AWS EC2 for training
- Validation is performed after per training epoch, and the best model weights is saved
- Use your own wandb API key in
- For inference, run
python ./inference.py
insemantic-segmentation/
- Preparation - Models
- Pretrained model by the authors on anime-specific images: Download
SketchColorizationModel.onnx
from Github Release, or use the following command to download from our Google Drive- Download:
gdown https://drive.google.com/u/2/uc?id=1G7nN4PL48TB6aJjCn7F6tJSM_d7pZ3_4
- Place into
colorization/
- Download:
- Pretrained model by the authors on anime-specific images: Download