# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/14c_segmask.core.ipynb.

# %% auto 0
__all__ = ['SegMask']

# %% ../../nbs/14c_segmask.core.ipynb 3
from fastcore.all import *
import cv2
import numpy as np
import polvo as pv

# %% ../../nbs/14c_segmask.core.ipynb 4
class SegMask:
    def __init__(self, mask, class_map):
        self.mask = mask
        self.class_map = class_map
       
    def accept_visit(self, visitor): return visitor.visit_segmask(self)
    
    def to_contours(self, method=cv2.CHAIN_APPROX_SIMPLE, mode=cv2.RETR_EXTERNAL):
        arr = self.mask.squeeze().astype(np.uint8)
        contours = {}
        for id_ in set(np.unique(arr)):
            label = self.class_map.id2name[id_]
            if label == 'background': continue
            mask = (arr == id_).astype(np.uint8)
            contours[label] = cv2.findContours(mask, mode, method)[0]
        return contours

    def __repr__(self):
        return f"<{self.__class__.__name__}\n{self.mask}>"
