from ..builder import PIPELINES
from ...utils.design_ensemble import get_seg_result


@PIPELINES.register_module()
class Segmentation(object):
    def __init__(self, device='cpu', show=False, debug=None):
        self.show = show
        self.device = device
        self.debug = debug

    def __call__(self, result):
        result['seg_result'] = get_seg_result(result["image_id"], result['image'])
        return result