imagemasking.py

# imagemasking.py by ALfuhrmann https://github.com/alicevision/meshroom/issues/188

from __future__ import print_function

__version__ = "1.0"

from meshroom.core import desc
import json
import numpy as np
import glob
import os
import cv2
import os.path
import logging

class MakeMaskFilter(desc.Node):

    inputs = [
        desc.File(
            name='input',
            label='Input',
            description='SfMData file.',
            value='',
            uid=[0],
        ),   
        desc.File(
            name="BackgroundImage",
            label='Background Image',
            description='Input filtered depth maps folder',
            value='',
            uid=[0],
        ),
        desc.ChoiceParam(
            name='blurSize',
            label='Blur Size',
            description='Size of blur kernel for noise reduction.',
            value=1,
            values=[0, 3, 5, 7, 9, 11, 13, 15],
            exclusive=True,
            uid=[0],
        ),
        desc.ChoiceParam(
            name='fixedRange',
            label='fill range',
            description='Flood fill value',
            value='fixed',
            values=['fixed', 'adaptive'],
            exclusive=True,
            uid=[0],
        ),
        desc.IntParam(
            name='fillThreshold',
            label='Fill tolerance',
            description='Tolerance threshold for background fill.',
            value=5,
            range=(0, 50, 1),
            uid=[0],
        ),
    ]

    outputs = [
        desc.File(
            name='output',
            label='Output',
            description='Output folder for mask images.',
            value=desc.Node.internalFolder,
            uid=[],
        ),
    ]

    displaysize = (1920, 850)
    kernel = np.ones((5, 5), np.uint8)

    def MakeMask(self, imgname, maskname):
        img = cv2.imread(imgname, cv2.IMREAD_GRAYSCALE)
        if self.doBlur:
            img = cv2.GaussianBlur(img, self.blurKernel, cv2.BORDER_DEFAULT)
        if self.subtractbkg:
            diff = cv2.absdiff(img, self.bkg)
        else:
            diff = img
        h, w = img.shape[:2]
        fillmask = np.zeros((h + 2, w + 2), np.uint8)  # mask = image + 1px border
        seeds = [(1, 1), (0, h - 1), (w - 1, h - 1), (w - 1, 0)]  # seedpoints in all corners
        for seedpoint in seeds:
            if fillmask[seedpoint[1] + 1, seedpoint[0] + 1] == 0:  # if not filled by a previous loop
                cv2.floodFill(diff, fillmask, seedpoint, 0, self.fillThreshold, self.fillThreshold,
                              flags=4 | (255 << 8) | self.fixedRange | cv2.FLOODFILL_MASK_ONLY)
        fillmask = 255 - fillmask[1:-1, 1:-1]  # remove border pixels again and invert values
        erosion = cv2.erode(fillmask, self.kernel, iterations=1)  # erode distance between mask and (blurred) input image
        # cv2.imshow("mask", cv2.resize(erosion / 255.0, self.displaysize))
        cv2.imwrite(maskname, erosion)
        # cv2.waitKey(1)

    def processChunk(self, chunk):
    
        print("Generate Mask Node start\n")

        self.doBlur = chunk.node.blurSize.value > 0
        self.blurKernel = (chunk.node.blurSize.value, chunk.node.blurSize.value)
        bkgname = chunk.node.BackgroundImage.value
        self.fillThreshold = chunk.node.fillThreshold.value

        self.subtractbkg = bkgname != ""
        if self.subtractbkg:
            self.bkg = cv2.imread(bkgname, cv2.IMREAD_GRAYSCALE)
            self.bkg = cv2.GaussianBlur(self.bkg, self.blurKernel, cv2.BORDER_DEFAULT)
        print("Blur kernel = {}".format(self.blurKernel))
        print("Background Image = '{}'".format(bkgname))
        print("Subtract Bkg = {}".format(self.subtractbkg))
        if chunk.node.fixedRange.value == "fixed":
            self.fixedRange = cv2.FLOODFILL_FIXED_RANGE
        else:
            self.fixedRange = 0
        with open(chunk.node.input.value, 'r') as f:
            sfm = json.load(f)
            for view in sfm["views"]:  # loop over all views/images
                viewId = view["viewId"]
                source_image_path = view["path"]

                inputfilename = source_image_path
                mask_name = os.path.join(chunk.node.output.value, viewId+".png")

                # masking depth maps
                print("input file: {}".format(inputfilename))
                print("mask  file: {}".format(mask_name))
                self.MakeMask(inputfilename, mask_name)

        print('Generate Mask Node end')