Diver's View Segmentation

Query image and prediction

BiWAKO.SUIMNet

Semantic segmentation model for diver's view dataset.

Attributes:

Name	Type	Description
model_path	str	Path to the onnx file. If automatic download is triggered, it is downloaded to this path.
session	onnxruntime.InferenceSession	Inference session.
h	int	Height of the input image.
w	int	Width of the input image.
input_name	str	Name of the input node.
output_name	str	Name of the output node.
c_map		Color map for the segmentation.

__init__(self, model='suim_rsb_72128', num_classes=5, **kwargs) special

Initialize the SUIMNet model.

Parameters:

Name	Type	Description	Default
model	str	Choice of the model or path to the downloaded onnx file. If the model hasn't been downloaded, it is automatically downloaded to this path. Defaults to "suim_rsb_72128".	'suim_rsb_72128'
num_classes	int	Number of classes to segmentate. Defaults to 5.	5

predict(self, image)

Return the segmentation map of the image.

Parameters:

Name	Type	Description	Default
image	Image	Image to segmentate. Accept path to the image or cv2 image.	required

Returns:

Type	Description
np.ndarray	Segmentation map of shape (h, w, num_classes). Each element is a confidence of the class.

render(self, prediction, image, **kwargs)

Apply the segmentation map to the image.

Parameters:

Name	Type	Description	Default
prediction	np.ndarray	Segmentation map of shape (h, w, num_classes) returned by predict().	required
image	Image	Input image. Accept path to the image or cv2 image.	required

Returns:

Type	Description
np.ndarray	Rendered image in cv2 format.