diff --git a/__init__.py b/__init__.py index 152a4cb..db9c5b7 100644 --- a/__init__.py +++ b/__init__.py @@ -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"] diff --git a/mat_anyone.py b/mat_anyone.py index 8dfd70d..9a6198d 100644 --- a/mat_anyone.py +++ b/mat_anyone.py @@ -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 diff --git a/run.py b/run.py index 3a63048..4f98548 100644 --- a/run.py +++ b/run.py @@ -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),) diff --git a/src/core/inference_core.py b/src/core/inference_core.py index af49913..daca51d 100644 --- a/src/core/inference_core.py +++ b/src/core/inference_core.py @@ -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 diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..a31aa4c --- /dev/null +++ b/utils.py @@ -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