Human Portrait Segmentation

Query image, predicted segmentation map, and visualization

BiWAKO.MODNet

Segmentation model trained on human portrait image.

Attributes:

Name	Type	Description
model	onnxruntime.InferenceSession	Inference session.
input_name	str	Name of input node.
output_name	str	Name of output node.
input_shape	tuple	Size of input image.
score_th	float	Threshold for mask.

__init__(self, model='modnet_256', score_th=0.5) special

Initialize MODNet.

Parameters:

Name	Type	Description	Default
model	str	Choice of model or path to the onnx file. Defaults to "modnet_256". If chosen model has not been downloaded, it will be downloaded automatically.	'modnet_256'
score_th	float	Optional threshold for mask used in self.render(). Any pixels in the mask with confidence score lower than this value will be set to 0. Defaults to 0.5.	0.5

_preprocess(self, image) private

Preprocess image for inference. This is automatically called by predict().

Preprocess

1. Resize image to the same size as the model input.
2. Normalize image to [-1, 1] with mean and std of 0.5.
3. Convert image to float32 and reshape to (1, C, H, W).

Parameters:

Name	Type	Description	Default
image	np.ndarray	Image in cv2 format.	required

Returns:

Type	Description
np.ndarray	Preprocessed image in numpy format.

predict(self, image)

Return mask of given image.

Parameters:

Name	Type	Description	Default
image	Image	Image to be segmented. Accept path to image or cv2 image.	required

Returns:

Type	Description
np.ndarray	Predicted mask in original size.

render(self, prediction, image, black_out=False, score_th=None, **kwargs)

Apply the mask to the input image.

Parameters:

Name	Type	Description	Default
prediction	np.ndarray	Mask returned by predict().	required
image	Image	Image to be segmented. Accept path to image or cv2 image.	required
black_out	bool	Whether to use black background. Defaults to False.	False
score_th	float	Optional threshold for mask. Defaults to use the value set in the constructor.	None

Returns:

Type	Description
np.ndarray	Segmented image in cv2 format.

Reference

https://github.com/ZHKKKe/MODNet/blob/master/onnx/inference_onnx.py