U2Net

Query image and prediction

BiWAKO.U2Net

Salient object segmentation model.

Attributes:

Name	Type	Description
IS	onnxruntime.InferenceSession	Inference session.
input_name	str	Input node name.
output_name	str	Output node name.
input_size	int	Input size. Set to 320.
mean	List[float]	Mean. Set to [0.485, 0.456, 0.406].
std	List[float]	Standard deviation. Set to [0.229, 0.224, 0.225].

__init__(self, model='mobile') special

U2Net Inference class.

Parameters:

Name	Type	Description	Default
model	str	Model name or downloaded onnx file. Accept one of ["basic", "mobile", "human_seg", "portrait"]. If model has not been downloaded, it will be downloaded automatically.	'mobile'

predict(self, image)

Return the predicted mask of the image.

Parameters:

Name	Type	Description	Default
image	Union[str, np.ndarray]	Image in cv2 format or path to image.	required

Returns:

Type	Description
np.ndarray	Predicted mask in cv2 format.

render(self, prediction, image)

Apply the predicted mask to the original image.

Parameters:

Name	Type	Description	Default
prediction	np.ndarray	Predicted mask in cv2 format.	required
image	Union[str, np.ndarray]	Image in cv2 format or path to image.	required

Returns:

Type	Description
np.ndarray	Rendered original image in cv2 format.