nodes/FuouM-ComfyUI-MatAnyone
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

finalizing

Browse Source
This commit is contained in:
FuouM
2025-03-24 09:36:05 +07:00
parent b940965c7d
commit ea8f9cca6d
5 changed files with 136 additions and 164 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+3 -1
View File
@@ -1,11 +1,13 @@
from .run import MatAnyoneVideo
from .run import MatAnyoneVideo, SolidColorBatched
NODE_CLASS_MAPPINGS = {
"MatAnyone": MatAnyoneVideo,
"SolidColorBatched": SolidColorBatched,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MatAnyone": "MatAnyone",
"SolidColorBatched": "Solid Color Batched",
}
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
mat_anyone.py
+22 -52
View File
@@ -44,6 +44,7 @@ def get_repeat(vframes: torch.Tensor, index: int, n_warmup: int):
def preprocess_mask(mask):
return mask * 255.0
def warming_up(mask, processor, n_warmup, repeated_frames):
for ti in tqdm(range(n_warmup), desc="Warming up"):
image = repeated_frames[ti]
@@ -58,56 +59,21 @@ def warming_up(mask, processor, n_warmup, repeated_frames):
return output_prob
# def inference_matanyone(
# vframes: torch.Tensor,
# mask: torch.Tensor,
# processor: InferenceCore,
# n_warmup=10,
# ):
# mask = preprocess_mask(mask) # from 0..1 to 0..255
# # repeat the first frame for warmup
# repeated_frames = get_repeat(vframes, 0, n_warmup)
# vframes = vframes[1:].to(device)
# length = vframes.shape[0]
# mask = mask.to(device)
# repeated_frames = repeated_frames.to(device)
# # inference start
# phas = []
# output_prob = warming_up(mask, processor, n_warmup, repeated_frames)
# # convert output probabilities to alpha matte
# # Get the last warm up frame
# mask = processor.output_prob_to_mask(output_prob)
# phas.append((mask).unsqueeze(0))
# # Process actual frames
# pbar = ProgressBar(length)
# for ti in tqdm(range(length), desc="Process frames"):
# image = vframes[ti]
# output_prob = processor.step(image)
# # convert output probabilities to alpha matte
# mask = processor.output_prob_to_mask(output_prob)
# phas.append((mask).unsqueeze(0))
# pbar.update_absolute(ti, length)
# return phas
def inference_matanyone(
vframes: torch.Tensor,
mask: torch.Tensor,
processor: InferenceCore,
frame_index: int=0,
frame_index: int = 0,
n_warmup=10,
):
mask = preprocess_mask(mask)
vframes = vframes.to(device)
mask = mask.to(device)
length = vframes.shape[0]
phas = [None] * length # initialize output list
phas = [None] * length # initialize output list
if frame_index >= length:
raise ValueError(f"Expected 0 <= x < {length}, got {frame_index}")
if frame_index == 0:
# Case 1: frame_index == 0 (original behavior)
@@ -124,13 +90,13 @@ def inference_matanyone(
phas[ti] = processor.output_prob_to_mask(output_prob).unsqueeze(0)
pbar.update_absolute(ti - 1, length - 1)
elif frame_index == length - 1:
# Case 2: frame_index == last frame (reverse propagation)
reversed_vframes = torch.flip(vframes, dims=[0])
repeated_frames = get_repeat(reversed_vframes, 0, n_warmup) # repeat the last frame
# repeat the last frame
repeated_frames = get_repeat(reversed_vframes, 0, n_warmup)
repeated_frames = repeated_frames.to(device)
reversed_mask = mask # mask for the last frame
reversed_mask = mask # mask for the last frame
output_prob = warming_up(reversed_mask, processor, n_warmup, repeated_frames)
phas[frame_index] = processor.output_prob_to_mask(output_prob).unsqueeze(0)
@@ -139,37 +105,41 @@ def inference_matanyone(
for ti in tqdm(range(1, length), desc="Backward Propagation"):
image = reversed_vframes[ti]
output_prob = processor.step(image)
phas[length - 1 - ti] = processor.output_prob_to_mask(output_prob).unsqueeze(0) # reverse index
phas[length - 1 - ti] = processor.output_prob_to_mask(
output_prob
).unsqueeze(0) # reverse index
pbar.update_absolute(ti - 1, length - 1)
elif 0 < frame_index < length - 1:
# Case 3: 0 < frame_index < last frame (forward and backward)
repeated_frames = get_repeat(vframes, frame_index, n_warmup) # repeat the frame_index frame
# repeat the frame_index frame
repeated_frames = get_repeat(vframes, frame_index, n_warmup)
repeated_frames = repeated_frames.to(device)
# Warm up at frame_index
output_prob = warming_up(mask, processor, n_warmup, repeated_frames)
phas[frame_index] = processor.output_prob_to_mask(output_prob).unsqueeze(0)
# Forward Propagation (from frame_index + 1 to end)
pbar_forward = ProgressBar(length - 1 - frame_index)
for ti in tqdm(range(frame_index + 1, length), desc="Forward Propagation"):
image = vframes[ti]
output_prob = processor.step(image)
phas[ti] = processor.output_prob_to_mask(output_prob).unsqueeze(0)
pbar_forward.update_absolute(ti - (frame_index + 1) , length - 1 - frame_index)
pbar_forward.update_absolute(
ti - (frame_index + 1), length - 1 - frame_index
)
# Backward Propagation (from frame_index - 1 to start)
reversed_vframes_backward = torch.flip(vframes[:frame_index], dims=[0]) # frames before frame_index, reversed
# frames before frame_index, reversed
reversed_vframes_backward = torch.flip(vframes[:frame_index], dims=[0])
pbar_backward = ProgressBar(frame_index)
for ti in tqdm(range(frame_index), desc="Backward Propagation"):
image = reversed_vframes_backward[ti]
output_prob = processor.step(image)
phas[frame_index - 1 - ti] = processor.output_prob_to_mask(output_prob).unsqueeze(0) # reverse index
phas[frame_index - 1 - ti] = processor.output_prob_to_mask(
output_prob
).unsqueeze(0) # reverse index
pbar_backward.update_absolute(ti, frame_index)
return phas
run.py
+72 -22
View File
@@ -7,6 +7,8 @@
from pathlib import Path
from .utils import get_mask, get_screen
base_dir = Path(__file__).resolve().parent
import torch
@@ -19,13 +21,6 @@ from .mat_anyone import (
from .src.core.inference_core import InferenceCore
def img_to_mask(tensor: torch.Tensor):
weights = torch.tensor([0.2989, 0.5870, 0.1140], device=tensor.device)
weights = weights.view(1, 3, 1, 1)
grayscale = torch.sum(tensor * weights, dim=1, keepdim=True)
return grayscale
class MatAnyoneVideo:
@classmethod
def INPUT_TYPES(s):
@@ -41,11 +36,21 @@ class MatAnyoneVideo:
{"default": 10, "min": 1, "step": 1},
),
},
"optional": {"foreground_mask": ("IMAGE",), "foreground_MASK": ("MASK",)},
"optional": {
"foreground_mask": ("IMAGE",),
"foreground_MASK": ("MASK",),
"solid_color": ("IMAGE",),
},
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("matte",)
RETURN_TYPES = (
"IMAGE",
"IMAGE",
)
RETURN_NAMES = (
"matte",
"green_screen",
)
FUNCTION = "todo"
CATEGORY = "MatAnyone"
@@ -56,6 +61,7 @@ class MatAnyoneVideo:
n_warmup: int,
foreground_mask: torch.Tensor | None = None,
foreground_MASK: torch.Tensor | None = None,
solid_color: torch.Tensor | None = None,
):
mask = get_mask(foreground_mask, foreground_MASK)
src_video = src_video.permute(0, 3, 1, 2) # T CHW RGB
@@ -65,20 +71,64 @@ class MatAnyoneVideo:
processor = InferenceCore(matanyone, cfg=matanyone.cfg)
phas = inference_matanyone(src_video, mask, processor, mask_frame, n_warmup)
out_mask = torch.cat(phas).unsqueeze(1).permute(0, 2, 3, 1)
out_mask_rgb = out_mask.repeat(1, 1, 1, 3) # Repeat the last dimension 3 times
gb = torch.empty(
0, out_mask_rgb.shape[1], out_mask_rgb.shape[2], out_mask_rgb.shape[3]
)
if solid_color is not None:
# T H W C
src_video_hwc = src_video.permute(0, 2, 3, 1).to(out_mask_rgb.device)
# x H W C -> T H W C , repeat batch to match video frames
solid_color_batched = solid_color.repeat(
src_video_hwc.shape[0], 1, 1, 1
).to(out_mask_rgb.device)
gb = out_mask_rgb * src_video_hwc + (1 - out_mask_rgb) * solid_color_batched
return (out_mask_rgb,)
return (
out_mask_rgb,
gb,
)
def get_mask(
foreground_mask: torch.Tensor | None = None,
foreground_MASK: torch.Tensor | None = None,
):
if foreground_mask is None and foreground_MASK is None:
raise ValueError("Please provide one mask image")
class SolidColorBatched:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"batch_size": (
"INT",
{"default": 1, "min": 1, "step": 1},
),
"height": (
"INT",
{"default": 720, "min": 1, "step": 1},
),
"width": (
"INT",
{"default": 1280, "min": 1, "step": 1},
),
"red": (
"INT",
{"default": 0, "min": 0, "step": 1},
),
"green": (
"INT",
{"default": 255, "min": 0, "step": 1},
),
"blue": (
"INT",
{"default": 0, "min": 0, "step": 1},
),
},
}
if foreground_MASK is not None:
mask = foreground_MASK.squeeze()
else:
mask = img_to_mask(foreground_mask.permute(0, 3, 1, 2)).squeeze()
return mask
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("solid",)
FUNCTION = "todo"
CATEGORY = "MatAnyone"
def todo(
self, batch_size: int, height: int, width: int, red: int, green: int, blue: int
):
return (get_screen(batch_size, height, width, red, green, blue),)
src/core/inference_core.py
+4 -89
View File
@@ -2,7 +2,6 @@ from typing import Iterable, List, Optional
import numpy as np
import torch
import torch.nn.functional as F
from .image_feature_store import ImageFeatureStore
from .memory_manager import MemoryManager
from .object_manager import ObjectManager
@@ -153,10 +152,7 @@ class InferenceCore:
Returns: (num_objects+1)*H*W normalized probability; the first channel is the background
"""
bs = key.shape[0]
if self.flip_aug:
assert bs == 2
else:
assert bs == 1
assert bs == 1
if not self.memory.engaged:
print("Trying to segment without any memory!")
@@ -199,13 +195,7 @@ class InferenceCore:
update_sensory=update_sensory,
)
# remove batch dim
if self.flip_aug:
# average predictions of the non-flipped and flipped version
pred_prob_with_bg = (
pred_prob_with_bg[0] + torch.flip(pred_prob_with_bg[1], dims=[-1])
) / 2
else:
pred_prob_with_bg = pred_prob_with_bg[0]
pred_prob_with_bg = pred_prob_with_bg[0]
if update_sensory:
self.memory.update_sensory(sensory, self.object_manager.all_obj_ids)
return pred_prob_with_bg
@@ -230,7 +220,6 @@ class InferenceCore:
mask: Optional[torch.Tensor] = None,
objects: Optional[List[int]] = None,
*,
idx_mask: bool = False,
end: bool = False,
delete_buffer: bool = True,
force_permanent: bool = False,
@@ -259,50 +248,12 @@ class InferenceCore:
force_permanent: the memory recorded this frame will be added to the permanent memory
"""
if objects is None and mask is not None:
assert not idx_mask
objects = list(range(1, mask.shape[0] + 1))
# resize input if needed -- currently only used for the GUI
resize_needed = False
if self.max_internal_size > 0:
h, w = image.shape[-2:]
min_side = min(h, w)
if min_side > self.max_internal_size:
resize_needed = True
new_h = int(h / min_side * self.max_internal_size)
new_w = int(w / min_side * self.max_internal_size)
image = F.interpolate(
image.unsqueeze(0),
size=(new_h, new_w),
mode="bilinear",
align_corners=False,
)[0]
if mask is not None:
if idx_mask:
mask = (
F.interpolate(
mask.unsqueeze(0).unsqueeze(0).float(),
size=(new_h, new_w),
mode="nearest-exact",
align_corners=False,
)[0, 0]
.round()
.long()
)
else:
mask = F.interpolate(
mask.unsqueeze(0),
size=(new_h, new_w),
mode="bilinear",
align_corners=False,
)[0]
self.curr_ti += 1
image, self.pad = pad_divide_by(image, 16) # DONE alreay for 3DCNN!!
image = image.unsqueeze(0) # add the batch dimension
if self.flip_aug:
image = torch.cat([image, torch.flip(image, dims=[-1])], dim=0)
# whether to update the working memory
is_mem_frame = (
@@ -348,17 +299,11 @@ class InferenceCore:
# merge predicted mask with the incomplete input mask
pred_prob_no_bg = pred_prob_with_bg[1:]
# use the mutual exclusivity of segmentation
if idx_mask:
pred_prob_no_bg[:, mask > 0] = 0
else:
pred_prob_no_bg[:, mask.max(0) > 0.5] = 0
pred_prob_no_bg[:, mask.max(0) > 0.5] = 0
new_masks = []
for mask_id, tmp_id in enumerate(corresponding_tmp_ids):
if idx_mask:
this_mask = (mask == objects[mask_id]).type_as(pred_prob_no_bg)
else:
this_mask = mask[tmp_id]
this_mask = mask[tmp_id]
if tmp_id > pred_prob_no_bg.shape[0]:
new_masks.append(this_mask.unsqueeze(0))
else:
@@ -366,24 +311,6 @@ class InferenceCore:
pred_prob_no_bg[tmp_id - 1] = this_mask
# new_masks are always in the order of tmp_id
mask = torch.cat([pred_prob_no_bg, *new_masks], dim=0)
elif idx_mask:
# simply convert cls to one-hot representation
if len(objects) == 0:
if delete_buffer:
self.image_feature_store.delete(self.curr_ti)
print("Trying to insert an empty mask as memory!")
return torch.zeros(
(1, key.shape[-2] * 16, key.shape[-1] * 16),
device=key.device,
dtype=key.dtype,
)
mask = torch.stack(
[
mask == objects[mask_id]
for mask_id, _ in enumerate(corresponding_tmp_ids)
],
dim=0,
)
if matting:
mask = mask.unsqueeze(0).float() / 255.0
pred_prob_with_bg = torch.cat([1 - mask, mask], 0)
@@ -392,10 +319,6 @@ class InferenceCore:
pred_prob_with_bg = torch.softmax(pred_prob_with_bg, dim=0)
self.last_mask = pred_prob_with_bg[1:].unsqueeze(0)
if self.flip_aug:
self.last_mask = torch.cat(
[self.last_mask, torch.flip(self.last_mask, dims=[-1])], dim=0
)
self.last_pix_feat = pix_feat
# save as memory if needed
@@ -429,14 +352,6 @@ class InferenceCore:
self.image_feature_store.delete(self.curr_ti)
output_prob = unpad(pred_prob_with_bg, self.pad)
if resize_needed:
# restore output to the original size
output_prob = F.interpolate(
output_prob.unsqueeze(0),
size=(h, w),
mode="bilinear",
align_corners=False,
)[0]
return output_prob
utils.py
+35
View File
@@ -0,0 +1,35 @@
import torch
def get_mask(
foreground_mask: torch.Tensor | None = None,
foreground_MASK: torch.Tensor | None = None,
):
if foreground_mask is None and foreground_MASK is None:
raise ValueError("Please provide one mask image")
if foreground_MASK is not None:
mask = foreground_MASK.squeeze()
else:
mask = img_to_mask(foreground_mask.permute(0, 3, 1, 2)).squeeze()
return mask
def img_to_mask(tensor: torch.Tensor):
weights = torch.tensor([0.2989, 0.5870, 0.1140], device=tensor.device)
weights = weights.view(1, 3, 1, 1)
grayscale = torch.sum(tensor * weights, dim=1, keepdim=True)
return grayscale
def get_screen(batch_size: int, height: int, width: int, r: int, g: int, b: int):
# Normalize RGB values to the range 0.0-1.0
r_normalized = float(r) / 255.0
g_normalized = float(g) / 255.0
b_normalized = float(b) / 255.0
rgb_image = torch.zeros((batch_size, height, width, 3), dtype=torch.float32)
rgb_image[:, :, :, 0] = r_normalized # Red channel (index 0)
rgb_image[:, :, :, 1] = g_normalized # Green channel (index 1)
rgb_image[:, :, :, 2] = b_normalized # Blue channel (index 2)
return rgb_image