From cbfb8bc0d2238a943158632aed0120bc36265903 Mon Sep 17 00:00:00 2001
From: spacepxl <spacetug93@gmail.com>
Date: Sat, 16 Dec 2023 19:35:37 -0500
Subject: [PATCH] Added alpha matte node

---
 nodes.py         | 79 +++++++++++++++++++++++++++++++++++++++++++++++-
 requirements.txt |  3 +-
 2 files changed, 80 insertions(+), 2 deletions(-)

diff --git a/nodes.py b/nodes.py
index 5d9b3ec..f511913 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1,11 +1,13 @@
 import torch
 import os
 import sys
+import copy
 
 import numpy as np
 import cv2
 from cv2.ximgproc import guidedFilter
-import copy
+from pymatting import *
+
 
 class AlphaClean:
     def __init__(self):
@@ -82,6 +84,79 @@ class AlphaClean:
         
         return (torch.from_numpy(i_dup),)
 
+class AlphaMatte:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "alpha_trimap": ("IMAGE",),
+                "preblur": ("INT", {
+                    "default": 8,
+                    "min": 0,
+                    "max": 256,
+                    "step": 1
+                }),
+                "blackpoint": ("FLOAT", {
+                    "default": 0.01,
+                    "min": 0.0,
+                    "max": 0.99,
+                    "step": 0.01
+                }),
+                "whitepoint": ("FLOAT", {
+                    "default": 0.99,
+                    "min": 0.01,
+                    "max": 1.0,
+                    "step": 0.01
+                }),
+                "max_iterations": ("INT", {
+                    "default": 1000,
+                    "min": 100,
+                    "max": 10000,
+                    "step": 100
+                }),
+                "estimate_fg": (["true", "false"],),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE",)
+    RETURN_NAMES = ("alpha", "fg", "bg",)
+    FUNCTION = "alpha_matte"
+
+    CATEGORY = "image/filters"
+
+    def alpha_matte(self, images, alpha_trimap, preblur, blackpoint, whitepoint, max_iterations, estimate_fg):
+        
+        d = preblur * 2 + 1
+        
+        i_dup = copy.deepcopy(images.cpu().numpy().astype(np.float64))
+        a_dup = copy.deepcopy(alpha_trimap.cpu().numpy().astype(np.float64))
+        fg = copy.deepcopy(images.cpu().numpy().astype(np.float64))
+        bg = copy.deepcopy(images.cpu().numpy().astype(np.float64))
+        
+        
+        for index, image in enumerate(i_dup):
+            trimap = a_dup[index][:,:,0] # convert to single channel
+            if preblur > 0:
+                trimap = cv2.GaussianBlur(trimap, (d, d), 0)
+            trimap = fix_trimap(trimap, blackpoint, whitepoint)
+            
+            alpha = estimate_alpha_cf(image, trimap, laplacian_kwargs={"epsilon": 1e-6}, cg_kwargs={"maxiter":max_iterations})
+            
+            if estimate_fg == "true":
+                fg[index], bg[index] = estimate_foreground_ml(image, alpha, return_background=True)
+            
+            a_dup[index] = np.stack([alpha, alpha, alpha], axis = -1) # convert back to rgb
+        
+        return (
+            torch.from_numpy(a_dup.astype(np.float32)), # alpha
+            torch.from_numpy(fg.astype(np.float32)), # fg
+            torch.from_numpy(bg.astype(np.float32)), # bg
+            )
+
 class BlurImageFast:
     def __init__(self):
         pass
@@ -367,6 +442,7 @@ class RemapRange:
 
 NODE_CLASS_MAPPINGS = {
     "AlphaClean": AlphaClean,
+    "AlphaMatte": AlphaMatte,
     "BlurImageFast": BlurImageFast,
     "BlurMaskFast": BlurMaskFast,
     "DilateErodeMask": DilateErodeMask,
@@ -377,6 +453,7 @@ NODE_CLASS_MAPPINGS = {
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "AlphaClean": "Alpha Clean",
+    "AlphaMatte": "Alpha Matte",
     "BlurImageFast": "Blur Image (Fast)",
     "BlurMaskFast": "Blur Mask (Fast)",
     "DilateErodeMask": "Dilate/Erode Mask",
diff --git a/requirements.txt b/requirements.txt
index 882fb05..629088f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1 +1,2 @@
-opencv-contrib-python
\ No newline at end of file
+opencv-contrib-python
+pymatting
\ No newline at end of file