{"id":13719463,"url":"https://github.com/Time0o/colorful-colorization","last_synced_at":"2025-05-07T11:32:01.833Z","repository":{"id":41469027,"uuid":"181223895","full_name":"Time0o/colorful-colorization","owner":"Time0o","description":"A from-scratch PyTorch implementation of  \"Colorful Image Colorization\" by Zhang et al.","archived":false,"fork":false,"pushed_at":"2021-06-16T18:49:06.000Z","size":26043,"stargazers_count":49,"open_issues_count":7,"forks_count":19,"subscribers_count":2,"default_branch":"master","last_synced_at":"2025-01-01T21:02:51.729Z","etag":null,"topics":["convolutional-neural-networks","deep-learning","image-colorization","pytorch"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/Time0o.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null}},"created_at":"2019-04-13T20:14:21.000Z","updated_at":"2024-07-31T07:07:53.000Z","dependencies_parsed_at":"2022-09-09T12:10:30.009Z","dependency_job_id":null,"html_url":"https://github.com/Time0o/colorful-colorization","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Time0o%2Fcolorful-colorization","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Time0o%2Fcolorful-colorization/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Time0o%2Fcolorful-colorization/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Time0o%2Fcolorful-colorization/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/Time0o","download_url":"https://codeload.github.com/Time0o/colorful-colorization/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":252869006,"owners_count":21816959,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["convolutional-neural-networks","deep-learning","image-colorization","pytorch"],"created_at":"2024-08-03T01:00:49.304Z","updated_at":"2025-05-07T11:32:01.207Z","avatar_url":"https://github.com/Time0o.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cp align=\"center\"\u003e\n  \u003cimg src=\"preview.png\"\u003e\n\u003c/p\u003e\n\n# Colorful Image Colorization PyTorch\n\nThis is a from-scratch PyTorch implementation of \"Colorful Image Colorization\"\n[1] by Zhang et al. created for the _Deep Learning in Data Science_ course at\nKTH Stockholm.\n\nThe following sections describe in detail:\n* how to install the dependencies necessary to get started with this project\n* how to colorize grayscale images using pretrained network weights\n* how to train the network on a new dataset\n* how to run colorization programatically\n\n## Prerequisites\n\nWe recommend you use Anaconda to create a virtual enviroment in which to\ninstall the modules needed to run this program, i.e you should run:\n\n```\nconda env create --file environment.yml\n```\n\nThere are some extra dependencies needed to run some of the scripts but you\nshould install those manually when it becomes necessary as you may not need\nthem.\n\nIn addition, you will need some files provided by R. Zhang, these include the\npoints of the ab gamut bins used to discretize the image labels and pretrained\nCaffe models. Run `resources/get_resources.sh` to download these automatically.\n**If you skip this step you will not be able to run the network at all, even if\nyou provide your own weight initialization or want to train from scratch**.\n\nYou might also notice another shell script, `data/get_cval.sh`, that downloads\nseveral additional resources. However, this is mainly a remnant of the\ndevelopement process and you can safely ignore it.\n\nIn order to use the pretrained weights for prediction, you will have to convert\nthem from Caffe to PyTorch. We provide the convenience script\n`scripts/convert_weights` for exactly this purpose. In order to use it you will\nhave to install the `caffe` Python module (if you want to convert one of the\nCaffe models provided by R. Zhang)\n\nFor example, in order to convert the Caffe model trained with class rebalancing\ndownloaded by `resources/get_resources.sh`, you can call the script like this:\n\n\n```\n./scripts/convert_weights vgg PYTORCH_WEIGHTS.tar \\\n\t--weights resources/colorization_release_v2.caffemodel \\\n\t--proto resources/colorization_deploy_v2.prototxt\n```\n\nWhich will save the converted PyTorch weights to `PYTORCH_WEIGHTS.tar`.\n\n## Colorize Images with Pretrained Weights\n\nThe easiest way to colorize several grayscale images of arbitrary size is to\nplace them in the same directory and colorize them in batch mode using\n`scripts/convert_images`. For example, if you have placed the images in\ndirectory `dir1` and subsequently run:\n\n```\n./scripts/convert_images predict_color \\\n    --input-dir dir1 \\\n    --output-dir dir2 \\\n    --model-checkpoint PYTORCH_WEIGHTS.tar \\\n    --gpu \\\n    --verbose\n```\n\nThe script will colorize all images in `dir1` on the GPU and place the results\nin `dir2` (with the same filenames). You can choose an annealed mean\ntemperature parameter other then the default 0.38 with `--annealed-mean-T`. .\n\n## Train the Network\n\n### Prepare a Dataset\n\nIf you intend to train the network on your own dataset, you might want to use\nthe convenience scripts `scripts/prepare_dataset` to convert it into a form\nsuitable for training. For example, if all your images are stored in a\ndirectory tree similar to this one:\n\n```\ndir1/\n├── subdir1\n│   ├── img1.JPEG\n│   ├── img2.JPEG\n│   └── ...\n├── subdir2\n│   ├── img1.JPEG\n│   ├── img2.JPEG\n│   └── ...\n└── ...\n\n```\n\nyou may want to run:\n\n```\n./scripts/prepare_dataset dir1 \\\n    --flatten \\\n    --purge \\\n    --clean \\\n    --file-ext JPEG \\\n    --val-split 0.2 \\\n    --test-split 0.1 \\\n    --resize-height 256 \\\n    --resize-width 256 \\\n    --verbose\n```\n\nThe script will first recursively look for images files with the extension\n`.JPEG` in `dir1` and remove all other files and those images that cannot be\nread or converted to RGB. It will then resize all remaining images to 256x256\nand randomly place them in the newly created subdirectories `train`, `val` and\n`test` using a 70/20/10 split.\n\nNote that this will take a while for large datasets since every single image\nhas to be read into memory. If your images already have the desired size (this\ndoes not necessarily have to be 256x256, the network is fully convolutional and\ncan train on images of arbitrary size) and you are sure that none of them are\ncorrupted, you don't have to use the `--resize-height/--resize-width` and\n`--clean` arguments which will speed up the process considerably.\n\n### Run the Training\n\nTo train the network on your dataset you can use the script\n`scripts/run_training`. The script accepts command line arguments that control\ne.g. the duration of the training and where/how often logfiles and model\ncheckpoints are written. More specific settings like dataloader configuration,\nnetwork type and optimizer settings need to be specified via a configuration\nfile which is essentially a nested directory of Python objects converted to\nJSON. Most likely you will want to use `config/default.json` and provide\nspecific settings or override some defaults in a separate JSON file. See\n`config/vgg.json` for an example.\n\nOnce you have decided on a configuration file you can run the script as follows:\n\n```\n./scripts/run_training \\\n    --config YOUR_CONFIG.json \\\n    --default-config config/default.json \\\n    --data-dir dir1 \\\n    --checkpoint-dir YOUR_CHECKPOINT_DIR \\\n    --log-file YOUR_LOG_FILE.txt \\\n    --iterations ITERATIONS \\\n    --iterations-till-checkpoint ITERATIONS_TILL_CHECKPOINT \\\n    --init-model-checkpoint INIT_MODEL-CHECKPOINT.tar\n```\n\nThis will recursively merge the configurations in `YOUR_CONFIG.json` and\n`config/default.json` and then train on the the images in `dir1` for\n`ITERATIONS` iterations (batches). Every `ITERATIONS_TILL_CHECKPOINT`\niterations, an intermediate model checkpoint will be written to\n`YOUR_CHECKPOINT_DIR`. Specifying `--init-model-checkpoint` is optional but\nuseful if you want to finetune the network from some pretrained set of weights.\n\nYou can also continue training from an arbitrary training checkpoint using the\n`--continue-training` flag which will load network weights and optimizer state\nfrom `INIT_MODEL_CHECKPOINT.tar` (which has to be a checkpoint created by a\nprevious run of `scripts/run_training`) and pick the training up from the last\ntraining iteration (thus `ITERATIONS` still specifies the total number of\ntraining iterations).\n\n## Colorize Images Programmatically\n\nColorizing images programmatically using our implementation is very simple. You\nfirst need to instantiate the network itself:\n\n```python\nfrom colorization import ColorizationNetwork\n\nnetwork = ColorizationNetwork(annealed_mean_T=0.38, device='gpu')\n```\n\nThe parameters should be self explanatory (and are in this case optional), use\n`device='cpu'` if you plan to run the network on the CPU.\n\nYou will then need to wrap the network in an instance of `ColorizationModel`\nwhich implements (among other things) checkpoint saving/loading:\n\n```python\nfrom colorization import ColorizationModel\n\nmodel = ColorizationModel(network)\nmodel.load_checkpoint('YOUR_CHECKPOINT_DIR/checkpoint_final.tar')\n```\n\nIn order to colorize a grayscale image you should then:\n* load it into a numpy array\n* resize a copy of it to 224x224 (this is not strictly necessary but produces\n  better results)\n* convert it to a torch tensor\n* pass it through the model\n* reassemble the result\n\nAll of this is already implemented in a convenience function:\n\n```python\nfrom colorization predict_color\nfrom skimage.io import imread\n\nimg = imread('YOUR_IMAGE.jpg')\nimg_colorized = predict_color(model, img)\n```\n\n## References\n\n[1] *Colorful Image Colorization*, Zhang, Richard and Isola, Phillip and Efros,\nAlexei A, in ECCV 2016\n([website](https://richzhang.github.io/colorization/))\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2FTime0o%2Fcolorful-colorization","html_url":"https://awesome.ecosyste.ms/projects/github.com%2FTime0o%2Fcolorful-colorization","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2FTime0o%2Fcolorful-colorization/lists"}