CascadePSP: Toward Class-Agnostic and Very High-Resolution Segmentation via Global and Local Refinement

Ho Kei Cheng*, Jihoon Chung*, Yu-Wing Tai, Chi-Keung Tang

Introduction

CascadePSP is a deep learning model for high-resolution segmentation refinement. This repository contains our PyTorch implementation with both training and testing functionalities. We also provide the annotated UHD dataset BIG and the pretrained model.

Here are some refinement results on high-resolution images.

Quick Start

Tested on PyTorch 1.0 -- though higher versions would likely work for inference as well.

Check out this folder. We have built a pip package that can refine an input image with two lines of code.

Install with

pip install segmentation-refinement

Code demo:

import cv2
import time
import matplotlib.pyplot as plt
import segmentation_refinement as refine
image = cv2.imread('test/aeroplane.jpg')
mask = cv2.imread('test/aeroplane.png', cv2.IMREAD_GRAYSCALE)

# model_path can also be specified here
# This step takes some time to load the model
refiner = refine.Refiner(device='cuda:0') # device can also be 'cpu'

# Fast - Global step only.
# Smaller L -> Less memory usage; faster in fast mode.
output = refiner.refine(image, mask, fast=False, L=900) 

# this line to save output
cv2.imwrite('output.png', output)

plt.imshow(output)
plt.show()

Network Overview

Global Step & Local Step

Global Step	Local Step

Refinement Module

Running:

Downloads:

More Results

Refining the masks of Human 3.6M

Image	Original Mask	Refined Mask

The first row is the failure case (see neck).

Credit

PSPNet implementation: https://github.com/Lextal/pspnet-pytorch

SyncBN implementation: https://github.com/vacancy/Synchronized-BatchNorm-PyTorch

If you find our work useful in your research, please cite the following:

@inproceedings{cheng2020cascadepsp,
  title={{CascadePSP}: Toward Class-Agnostic and Very High-Resolution Segmentation via Global and Local Refinement},
  author={Cheng, Ho Kei and Chung, Jihoon and Tai, Yu-Wing and Tang, Chi-Keung},
  booktitle={CVPR},
  year={2020}
}

lyclyc52/CascadePSP