Simple inference implementation with trained HRNet on MIT ADE20K dataset, using PyTorch 1.6.0. Most of the code taken from [1]. Purpose of this project is to unify sky pixels with ultra high prediction confidence to a single color, in order to remove sun light effect and inconsistent cloud texture.
-
Download pretrained model from http://sceneparsing.csail.mit.edu/model/pytorch and store them in
./ade20k-hrnetv2-c1/
. -
Specify your test input image directory and test output image directory in
semantic_segmentation.py
.
image_path = './input/'
output_path = './output/'
- Specify if you want to use GPU or not (-1 for cpu, others for gpu index) in
semantic_segmentation.py
.
gpu = -1
-
Run
python semantic_segmentation.py
-
All the inference results will be stored in your output_path
[1] Original code: https://github.com/CSAILVision/semantic-segmentation-pytorch.
[2] HRNet: https://arxiv.org/abs/1904.04514.