Source code for node_cv_disparity

# This file is part of Sympathy for Data.
# Copyright (c) 2018, Combine Control Systems AB
#
# SYMPATHY FOR DATA COMMERCIAL LICENSE
# You should have received a link to the License with Sympathy for Data.

from sympathy.api import node, exceptions
from sympathy.api.nodeconfig import Ports, Tag, Tags

import numpy as np

from sylib.imageprocessing.image import ImagePort
from sympathy.utils.pip_util import import_optional


def _cv2():
    return import_optional("cv2", group="opencv")




class ComputeDisparity(node.Node):

    author = 'Mathias Broxvall'
    icon = 'image_depthdisparity.svg'
    description = (
        'Creates a disparity map from a stereo image (two images with an\n'
        'offset along the x-axis). This disparity map shows how much\n'
        'each block of pixels in the the first image needs to be\n'
        'shifted in order to map to the closest block in the other\n'
        'image. Can be used as a pseudo-depth map.')
    tags = Tags(Tag.ImageProcessing.ImageManipulation)
    name = 'Compute stereo disparity (OpenCV)'
    nodeid = 'com.sympathyfordata.imageanalysis.disparity'

    parameters = node.parameters()
    parameters.set_integer(
        'max_disparity', value=64, label='Max disparity',
        description=(
            'Largest possible disparity (in pixels), also discards this many\n'
            'pixels from the edge of the image since they cannot match a\n'
            'corresponding pixel in the other image')
    )
    parameters.set_integer(
        'block_size', value=5, label='Block size',
        description=(
            'Size of the blocks compared between the two images. Larger '
            'values\n'
            'give better matches but a coarser resolution of the resulting\n'
            'disparity map. This number must be positive and odd.'),
    )

    __doc__ = description

    inputs = Ports([
        ImagePort('Left image', name='left'),
        ImagePort('Right image', name='right'),
    ])
    outputs = Ports([
        ImagePort('Result', name='result'),
    ])

    def execute(self, node_context):
        params = node_context.parameters
        left = node_context.input['left'].get_image()
        right = node_context.input['right'].get_image()
        result = node_context.output['result']
        max_disparity = params['max_disparity'].value
        block_size = params['block_size'].value

        if int(left.sum()) == 0 or int(right.sum()) == 0:
            raise exceptions.SyDataError("Empty table")

        if np.max(left) < 1.0:
            left = left * 255
        if np.max(right) < 1.0:
            right = right * 255
        left = left[:, :, 0].astype(np.uint8)
        right = right[:, :, 0].astype(np.uint8)

        stereo_mapper = _cv2().StereoBM_create(max_disparity, block_size)
        disparity = stereo_mapper.compute(left, right)
        result.set_image(disparity)