From 56d194a4300610a17cc3c6682aaf63f039e9c4ab Mon Sep 17 00:00:00 2001
From: wangzhibo <wangzhibo@shsening.com>
Date: 星期日, 03 八月 2025 16:11:37 +0800
Subject: [PATCH] add matanyone-better auqliaty than robustvideomatting
---
auto_media_publisher/modifiers/video_matting/matanyone/model/utils/resnet.py | 179 ++
auto_media_publisher/modifiers/video_matting/robustvideomatting/lraspp.py | 0
auto_media_publisher/modifiers/video_matting/matanyone/inference/object_info.py | 24
auto_media_publisher/modifiers/video_matting/matanyone/utils/tensor_utils.py | 62
auto_media_publisher/modifiers/video_matting/robust_video_matting_converter.py | 22
auto_media_publisher/modifiers/video_matting/matanyone/model/__init__.py | 0
auto_media_publisher/utils/media_input_output.py | 160 +
auto_media_publisher/modifiers/video_matting/robustvideomatting/decoder.py | 0
auto_media_publisher/utils/utils_yolo.py | 79 +
auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/positional_encoding.py | 108 +
auto_media_publisher/modifiers/video_matting/robustvideomatting/__init__.py | 0
data/auto_media_publisher.sqlite3 | 0
auto_media_publisher/modifiers/video_matting/matanyone/inference/inference_core.py | 609 +++++++
auto_media_publisher/modifiers/video_matting/matanyone/__init__.py | 0
auto_media_publisher/modifiers/video_matting/matanyone/inference/kv_memory_store.py | 348 ++++
auto_media_publisher/modifiers/video_matting/matanyone/model/matanyone.py | 333 ++++
auto_media_publisher/generators/docx_generator_gemini.py | 4
auto_media_publisher/modifiers/video_matting/robustvideomatting/resnet.py | 0
auto_media_publisher/utils/utils_text2image_pollinations.py | 31
auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/transformer_layers.py | 161 ++
auto_media_publisher/modifiers/video_matting/mat_anyone.py | 51
auto_media_publisher/modifiers/video_matting/matanyone/inference/image_feature_store.py | 56
auto_media_publisher/modifiers/video_matting/matanyone/inference/object_manager.py | 149 +
auto_media_publisher/modifiers/video_matting/matanyone/inference/utils/__init__.py | 0
auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/object_summarizer.py | 89 +
auto_media_publisher/modifiers/video_matting/matanyone/model/utils/memory_utils.py | 107 +
test/test_text2image.py | 2
auto_media_publisher/modifiers/rewriter/docx_rewriter_gemini.py | 2
auto_media_publisher/modifiers/video_matting/robustvideomatting/fast_guided_filter.py | 0
auto_media_publisher/utils/utils_images.py | 56
auto_media_publisher/modifiers/video_matting/matanyone/model/utils/parameter_groups.py | 72
auto_media_publisher/modifiers/video_matting/robustvideomatting/model.py | 0
auto_media_publisher/modifiers/video_matting/matanyone/model/channel_attn.py | 39
test/test_video_matting.py | 8
auto_media_publisher/modifiers/video_matting/matanyone/utils/get_default_model.py | 27
auto_media_publisher/modifiers/video_matting/matanyone/model/big_modules.py | 365 ++++
.gitignore | 2
auto_media_publisher/modifiers/video_matting/matanyone/model/utils/__init__.py | 0
auto_media_publisher/modifiers/video_matting/matanyone/inference/utils/args_utils.py | 30
auto_media_publisher/modifiers/video_matting/robustvideomatting/mobilenetv3.py | 0
auto_media_publisher/modifiers/video_matting/matanyone/utils/inference_utils.py | 57
data/cookies/18918594752/gongzhonghao.json | 2
scripts/work_gen_images.py | 46
scripts/girls.json | 99 +
scripts/office.json | 279 +++
auto_media_publisher/modifiers/video_matting/matanyone/inference/memory_manager.py | 453 +++++
auto_media_publisher/modifiers/video_matting/robustvideomatting/deep_guided_filter.py | 0
auto_media_publisher/modifiers/video_matting/matanyone/model/modules.py | 149 +
auto_media_publisher/modifiers/video_matting/matanyone/inference/__init__.py | 0
auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/object_transformer.py | 206 ++
auto_media_publisher/modifiers/video_matting/matanyone/model/aux_modules.py | 93 +
/dev/null | 0
test/test_face_swap.py | 2
auto_media_publisher/modifiers/video_matting/matanyone/model/group_modules.py | 126 +
auto_media_publisher/modifiers/video_matting/matanyone/utils/__init__.py | 0
scripts/__init__.py | 0
test/test_matanyone.py | 25
auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/__init__.py | 0
58 files changed, 4,668 insertions(+), 44 deletions(-)
diff --git a/.gitignore b/.gitignore
index 202bfc8..10428fe 100644
--- a/.gitignore
+++ b/.gitignore
@@ -32,3 +32,5 @@
# 鍏朵粬涓嶅お鎯崇鍏宠仈鐨�
涓嬭浇19妤艰鍧�/
data/cookies/18918594752/gongzhonghao.json
+data/auto_media_publisher.sqlite3
+scripts/output/
diff --git a/auto_media_publisher/generators/docx_generator_gemini.py b/auto_media_publisher/generators/docx_generator_gemini.py
index 292036b..e9f076d 100644
--- a/auto_media_publisher/generators/docx_generator_gemini.py
+++ b/auto_media_publisher/generators/docx_generator_gemini.py
@@ -88,7 +88,7 @@
self._add_content_to_doc(doc,"text",para)
if random.random() > 0.5:
try:
- img_stream = self.text2imager.generate_image_to_stream(para)
+ img_stream = self.text2imager.generate_image_with_trans_to_stream(para)
self._add_content_to_doc(doc,"image",img_stream)
except:
pass
@@ -99,7 +99,7 @@
self.summary_info = self._summarize_title_intro_keywords()
print(self.summary_info)
- self.text2imager.generate_image(self.summary_info["description"],output_path.replace(".docx",".png"))
+ self.text2imager.generate_image_with_trans(self.summary_info["description"],output_path.replace(".docx",".png"))
self.summary_info["path_cover"] = output_path.replace(".docx",".png")
diff --git a/auto_media_publisher/modifiers/rewriter/docx_rewriter_gemini.py b/auto_media_publisher/modifiers/rewriter/docx_rewriter_gemini.py
index 09c91b9..cc8e607 100644
--- a/auto_media_publisher/modifiers/rewriter/docx_rewriter_gemini.py
+++ b/auto_media_publisher/modifiers/rewriter/docx_rewriter_gemini.py
@@ -64,7 +64,7 @@
self.summary_info = self._summarize_title_intro_keywords()
print(self.summary_info)
- self.text2imager.generate_image(self.summary_info["description"],output_path.replace(".docx",".png"))
+ self.text2imager.generate_image_with_trans(self.summary_info["description"],output_path.replace(".docx",".png"))
self.summary_info["path_cover"] = output_path.replace(".docx",".png")
diff --git a/auto_media_publisher/modifiers/video_matting/mat_anyone.py b/auto_media_publisher/modifiers/video_matting/mat_anyone.py
new file mode 100644
index 0000000..ebc7e6c
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/mat_anyone.py
@@ -0,0 +1,51 @@
+from auto_media_publisher.modifiers.video_matting.matanyone.inference.inference_core import InferenceCore
+from auto_media_publisher.utils.utils_yolo import generate_mask_from_first_frame
+from auto_media_publisher.utils.logger import get_logger
+import os
+import traceback
+MAT_ANYONE_MODEL_DIR ="./models/MatAnyone"
+
+
+class MatAnyoneConverter:
+ """
+ 瑙嗛鎶犲儚杞崲鍣�-- 鍙浆鎹㈠浜轰腑鐨勪竴浜恒��
+ 灏佽鍘熸潵鐨�
+ """
+
+ def __init__(self):
+ self.logger = get_logger('MatAnyoneConverter','admin')
+ self.processor = InferenceCore(MAT_ANYONE_MODEL_DIR)
+
+
+ def video_matting(self,input_path,output_path,bg_video=None,bg_image=None,filter_type="dist"):
+ if not input_path or (input_path and not os.path.exists(input_path)):
+ print(f"input video is NOT exist.{input_path}")
+ return False
+ if bg_video and not os.path.exists(bg_video):
+ print(f"input background video is NOT exist.{bg_video}")
+ return False
+ if bg_image and not os.path.exists(bg_image):
+ print(f"input background image is NOT exist.{bg_image}")
+ return False
+
+ try:
+ mask_ = generate_mask_from_first_frame(video_path=input_path,filter_type=filter_type,mask_type="file")
+
+ if mask_:
+
+ foreground_path, alpha_path = self.processor.process_video(
+ input_path = input_path,
+ mask_path = mask_,
+ output_path = output_path,
+ bg_video = bg_video,
+ bg_image = bg_image
+ )
+ return True
+ return False
+
+ except Exception as e:
+ traceback.extract_stack()
+ self.logger.error(f"[ error ], 鍙戠敓閿欒: {e}")
+ return False
+
+
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/__init__.py b/auto_media_publisher/modifiers/video_matting/matanyone/__init__.py
similarity index 100%
copy from auto_media_publisher/modifiers/video_matting/matting_model/__init__.py
copy to auto_media_publisher/modifiers/video_matting/matanyone/__init__.py
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/__init__.py b/auto_media_publisher/modifiers/video_matting/matanyone/inference/__init__.py
similarity index 100%
copy from auto_media_publisher/modifiers/video_matting/matting_model/__init__.py
copy to auto_media_publisher/modifiers/video_matting/matanyone/inference/__init__.py
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/inference/image_feature_store.py b/auto_media_publisher/modifiers/video_matting/matanyone/inference/image_feature_store.py
new file mode 100644
index 0000000..824fa7a
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/inference/image_feature_store.py
@@ -0,0 +1,56 @@
+import warnings
+from typing import Iterable
+import torch
+from auto_media_publisher.modifiers.video_matting.matanyone.model.matanyone import MatAnyone
+
+
+class ImageFeatureStore:
+ """
+ A cache for image features.
+ These features might be reused at different parts of the inference pipeline.
+ This class provide an interface for reusing these features.
+ It is the user's responsibility to delete redundant features.
+
+ Feature of a frame should be associated with a unique index -- typically the frame id.
+ """
+ def __init__(self, network: MatAnyone, no_warning: bool = False):
+ self.network = network
+ self._store = {}
+ self.no_warning = no_warning
+
+ def _encode_feature(self, index: int, image: torch.Tensor, last_feats=None) -> None:
+ ms_features, pix_feat = self.network.encode_image(image, last_feats=last_feats)
+ key, shrinkage, selection = self.network.transform_key(ms_features[0])
+ self._store[index] = (ms_features, pix_feat, key, shrinkage, selection)
+
+ def get_all_features(self, images: torch.Tensor) -> (Iterable[torch.Tensor], torch.Tensor):
+ seq_length = images.shape[0]
+ ms_features, pix_feat = self.network.encode_image(images, seq_length)
+ key, shrinkage, selection = self.network.transform_key(ms_features[0])
+ for index in range(seq_length):
+ self._store[index] = ([f[index].unsqueeze(0) for f in ms_features], pix_feat[index].unsqueeze(0), key[index].unsqueeze(0), shrinkage[index].unsqueeze(0), selection[index].unsqueeze(0))
+
+ def get_features(self, index: int,
+ image: torch.Tensor, last_feats=None) -> (Iterable[torch.Tensor], torch.Tensor):
+ if index not in self._store:
+ self._encode_feature(index, image, last_feats)
+
+ return self._store[index][:2]
+
+ def get_key(self, index: int,
+ image: torch.Tensor, last_feats=None) -> (torch.Tensor, torch.Tensor, torch.Tensor):
+ if index not in self._store:
+ self._encode_feature(index, image, last_feats)
+
+ return self._store[index][2:]
+
+ def delete(self, index: int) -> None:
+ if index in self._store:
+ del self._store[index]
+
+ def __len__(self):
+ return len(self._store)
+
+ def __del__(self):
+ if len(self._store) > 0 and not self.no_warning:
+ warnings.warn(f'Leaking {self._store.keys()} in the image feature store')
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/inference/inference_core.py b/auto_media_publisher/modifiers/video_matting/matanyone/inference/inference_core.py
new file mode 100644
index 0000000..98eccac
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/inference/inference_core.py
@@ -0,0 +1,609 @@
+import logging
+from omegaconf import DictConfig
+from typing import List, Optional, Iterable, Union,Tuple
+
+import os
+import cv2
+import torch
+import imageio
+import tempfile
+import numpy as np
+from tqdm import tqdm
+from PIL import Image
+import torch.nn.functional as F
+from torchvision import transforms
+
+from auto_media_publisher.modifiers.video_matting.matanyone.inference.memory_manager import MemoryManager
+from auto_media_publisher.modifiers.video_matting.matanyone.inference.object_manager import ObjectManager
+from auto_media_publisher.modifiers.video_matting.matanyone.inference.image_feature_store import ImageFeatureStore
+from auto_media_publisher.modifiers.video_matting.matanyone.model.matanyone import MatAnyone
+from auto_media_publisher.modifiers.video_matting.matanyone.utils.tensor_utils import pad_divide_by, unpad, aggregate
+from auto_media_publisher.modifiers.video_matting.matanyone.utils.inference_utils import gen_dilate, gen_erosion, read_frame_from_videos
+from auto_media_publisher.utils.media_input_output import VideoWriter,VideoReader
+log = logging.getLogger()
+
+
+class InferenceCore:
+
+ def __init__(self,
+ network: Union[MatAnyone,str],
+ cfg: DictConfig = None,
+ *,
+ image_feature_store: ImageFeatureStore = None,
+ device: Union[str, torch.device] = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ ):
+ if isinstance(network, str):
+ #network = MatAnyone.from_pretrained(network)
+ network = MatAnyone.from_pretrained(network, local_files_only=True)
+ network.to(device)
+ network.eval()
+ self.device =device
+ self.network = network
+ cfg = cfg if cfg is not None else network.cfg
+ self.cfg = cfg
+ self.mem_every = cfg.mem_every
+ stagger_updates = cfg.stagger_updates
+ self.chunk_size = cfg.chunk_size
+ self.save_aux = cfg.save_aux
+ self.max_internal_size = cfg.max_internal_size
+ self.flip_aug = cfg.flip_aug
+
+ self.curr_ti = -1
+ self.last_mem_ti = 0
+ # at which time indices should we update the sensory memory
+ if stagger_updates >= self.mem_every:
+ self.stagger_ti = set(range(1, self.mem_every + 1))
+ else:
+ self.stagger_ti = set(
+ np.round(np.linspace(1, self.mem_every, stagger_updates)).astype(int))
+ self.object_manager = ObjectManager()
+ self.memory = MemoryManager(cfg=cfg, object_manager=self.object_manager)
+
+ if image_feature_store is None:
+ self.image_feature_store = ImageFeatureStore(self.network)
+ else:
+ self.image_feature_store = image_feature_store
+
+ self.last_mask = None
+ self.last_pix_feat = None
+ self.last_msk_value = None
+
+ def clear_memory(self):
+ self.curr_ti = -1
+ self.last_mem_ti = 0
+ self.memory = MemoryManager(cfg=self.cfg, object_manager=self.object_manager)
+
+ def clear_non_permanent_memory(self):
+ self.curr_ti = -1
+ self.last_mem_ti = 0
+ self.memory.clear_non_permanent_memory()
+
+ def clear_sensory_memory(self):
+ self.curr_ti = -1
+ self.last_mem_ti = 0
+ self.memory.clear_sensory_memory()
+
+ def update_config(self, cfg):
+ self.mem_every = cfg['mem_every']
+ self.memory.update_config(cfg)
+
+ def clear_temp_mem(self):
+ self.memory.clear_work_mem()
+ # self.object_manager = ObjectManager()
+ self.memory.clear_obj_mem()
+ # self.memory.clear_sensory_memory()
+
+ def _add_memory(self,
+ image: torch.Tensor,
+ pix_feat: torch.Tensor,
+ prob: torch.Tensor,
+ key: torch.Tensor,
+ shrinkage: torch.Tensor,
+ selection: torch.Tensor,
+ *,
+ is_deep_update: bool = True,
+ force_permanent: bool = False) -> None:
+ """
+ Memorize the given segmentation in all memory stores.
+
+ The batch dimension is 1 if flip augmentation is not used.
+ image: RGB image, (1/2)*3*H*W
+ pix_feat: from the key encoder, (1/2)*_*H*W
+ prob: (1/2)*num_objects*H*W, in [0, 1]
+ key/shrinkage/selection: for anisotropic l2, (1/2)*_*H*W
+ selection can be None if not using long-term memory
+ is_deep_update: whether to use deep update (e.g. with the mask encoder)
+ force_permanent: whether to force the memory to be permanent
+ """
+ if prob.shape[1] == 0:
+ # nothing to add
+ log.warn('Trying to add an empty object mask to memory!')
+ return
+
+ if force_permanent:
+ as_permanent = 'all'
+ else:
+ as_permanent = 'first'
+
+ self.memory.initialize_sensory_if_needed(key, self.object_manager.all_obj_ids)
+ msk_value, sensory, obj_value, _ = self.network.encode_mask(
+ image,
+ pix_feat,
+ self.memory.get_sensory(self.object_manager.all_obj_ids),
+ prob,
+ deep_update=is_deep_update,
+ chunk_size=self.chunk_size,
+ need_weights=self.save_aux)
+ self.memory.add_memory(key,
+ shrinkage,
+ msk_value,
+ obj_value,
+ self.object_manager.all_obj_ids,
+ selection=selection,
+ as_permanent=as_permanent)
+ self.last_mem_ti = self.curr_ti
+ if is_deep_update:
+ self.memory.update_sensory(sensory, self.object_manager.all_obj_ids)
+ self.last_msk_value = msk_value
+
+ def _segment(self,
+ key: torch.Tensor,
+ selection: torch.Tensor,
+ pix_feat: torch.Tensor,
+ ms_features: Iterable[torch.Tensor],
+ update_sensory: bool = True) -> torch.Tensor:
+ """
+ Produce a segmentation using the given features and the memory
+
+ The batch dimension is 1 if flip augmentation is not used.
+ key/selection: for anisotropic l2: (1/2) * _ * H * W
+ pix_feat: from the key encoder, (1/2) * _ * H * W
+ ms_features: an iterable of multiscale features from the encoder, each is (1/2)*_*H*W
+ with strides 16, 8, and 4 respectively
+ update_sensory: whether to update the sensory memory
+
+ 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
+
+ if not self.memory.engaged:
+ log.warn('Trying to segment without any memory!')
+ return torch.zeros((1, key.shape[-2] * 16, key.shape[-1] * 16),
+ device=key.device,
+ dtype=key.dtype)
+
+ uncert_output = None
+
+ if self.curr_ti == 0: # ONLY for the first frame for prediction
+ memory_readout = self.memory.read_first_frame(self.last_msk_value, pix_feat, self.last_mask, self.network, uncert_output=uncert_output)
+ else:
+ memory_readout = self.memory.read(pix_feat, key, selection, self.last_mask, self.network, uncert_output=uncert_output, last_msk_value=self.last_msk_value, ti=self.curr_ti,
+ last_pix_feat=self.last_pix_feat, last_pred_mask=self.last_mask)
+ memory_readout = self.object_manager.realize_dict(memory_readout)
+
+ sensory, _, pred_prob_with_bg = self.network.segment(ms_features,
+ memory_readout,
+ self.memory.get_sensory(
+ self.object_manager.all_obj_ids),
+ chunk_size=self.chunk_size,
+ 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]
+ if update_sensory:
+ self.memory.update_sensory(sensory, self.object_manager.all_obj_ids)
+ return pred_prob_with_bg
+
+ def pred_all_flow(self, images):
+ self.total_len = images.shape[0]
+ images, self.pad = pad_divide_by(images, 16)
+ images = images.unsqueeze(0) # add the batch dimension: (1,t,c,h,w)
+
+ self.flows_forward, self.flows_backward = self.network.pred_forward_backward_flow(images)
+
+ def encode_all_images(self, images):
+ images, self.pad = pad_divide_by(images, 16)
+ self.image_feature_store.get_all_features(images) # t c h w
+ return images
+
+ def step(self,
+ image: torch.Tensor,
+ 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,
+ matting: bool = True,
+ first_frame_pred: bool = False) -> torch.Tensor:
+ """
+ Take a step with a new incoming image.
+ If there is an incoming mask with new objects, we will memorize them.
+ If there is no incoming mask, we will segment the image using the memory.
+ In both cases, we will update the memory and return a segmentation.
+
+ image: 3*H*W
+ mask: H*W (if idx mask) or len(objects)*H*W or None
+ objects: list of object ids that are valid in the mask Tensor.
+ The ids themselves do not need to be consecutive/in order, but they need to be
+ in the same position in the list as the corresponding mask
+ in the tensor in non-idx-mask mode.
+ objects is ignored if the mask is None.
+ If idx_mask is False and objects is None, we sequentially infer the object ids.
+ idx_mask: if True, mask is expected to contain an object id at every pixel.
+ If False, mask should have multiple channels with each channel representing one object.
+ end: if we are at the end of the sequence, we do not need to update memory
+ if unsure just set it to False
+ delete_buffer: whether to delete the image feature buffer after this step
+ 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 = ((self.curr_ti - self.last_mem_ti >= self.mem_every) or
+ (mask is not None)) and (not end)
+ # segment when there is no input mask or when the input mask is incomplete
+ need_segment = (mask is None) or (self.object_manager.num_obj > 0
+ and not self.object_manager.has_all(objects))
+ update_sensory = ((self.curr_ti - self.last_mem_ti) in self.stagger_ti) and (not end)
+
+ # reinit if it is the first frame for prediction
+ if first_frame_pred:
+ self.curr_ti = 0
+ self.last_mem_ti = 0
+ is_mem_frame = True
+ need_segment = True
+ update_sensory = True
+
+ # encoding the image
+ ms_feat, pix_feat = self.image_feature_store.get_features(self.curr_ti, image)
+ key, shrinkage, selection = self.image_feature_store.get_key(self.curr_ti, image)
+
+ # segmentation from memory if needed
+ if need_segment:
+ pred_prob_with_bg = self._segment(key,
+ selection,
+ pix_feat,
+ ms_feat,
+ update_sensory=update_sensory)
+
+ # use the input mask if provided
+ if mask is not None:
+ # inform the manager of the new objects, and get a list of temporary id
+ # temporary ids -- indicates the position of objects in the tensor
+ # (starts with 1 due to the background channel)
+ corresponding_tmp_ids, _ = self.object_manager.add_new_objects(objects)
+
+ mask, _ = pad_divide_by(mask, 16)
+ if need_segment:
+ # 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
+
+ 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]
+ if tmp_id > pred_prob_no_bg.shape[0]:
+ new_masks.append(this_mask.unsqueeze(0))
+ else:
+ # +1 for padding the background channel
+ 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)
+ log.warn('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.
+ pred_prob_with_bg = torch.cat([1-mask, mask], 0)
+ else:
+ pred_prob_with_bg = aggregate(mask, dim=0)
+ 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
+ if is_mem_frame or force_permanent:
+ # clear the memory for given mask and add the first predicted mask
+ if first_frame_pred:
+ self.clear_temp_mem()
+ self._add_memory(image,
+ pix_feat,
+ self.last_mask,
+ key,
+ shrinkage,
+ selection,
+ force_permanent=force_permanent,
+ is_deep_update=True)
+ else: # compute self.last_msk_value for non-memory frame
+ msk_value, _, _, _ = self.network.encode_mask(
+ image,
+ pix_feat,
+ self.memory.get_sensory(self.object_manager.all_obj_ids),
+ self.last_mask,
+ deep_update=False,
+ chunk_size=self.chunk_size,
+ need_weights=self.save_aux)
+ self.last_msk_value = msk_value
+
+ if delete_buffer:
+ 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
+
+ def delete_objects(self, objects: List[int]) -> None:
+ """
+ Delete the given objects from the memory.
+ """
+ self.object_manager.delete_objects(objects)
+ self.memory.purge_except(self.object_manager.all_obj_ids)
+
+ def output_prob_to_mask(self, output_prob: torch.Tensor, matting: bool = True) -> torch.Tensor:
+ if matting:
+ new_mask = output_prob[1:].squeeze(0)
+ else:
+ mask = torch.argmax(output_prob, dim=0)
+
+ # index in tensor != object id -- remap the ids here
+ new_mask = torch.zeros_like(mask)
+ for tmp_id, obj in self.object_manager.tmp_id_to_obj.items():
+ new_mask[mask == tmp_id] = obj.id
+
+ return new_mask
+
+ @torch.inference_mode()
+ @torch.amp.autocast("cuda")
+ def process_video(
+ self,
+ input_path: str,
+ mask_path: str,
+ output_path: str = None,
+ bg_video :str = None,
+ bg_image : str = None,
+ n_warmup: int = 10,
+ r_erode: int = 10,
+ r_dilate: int = 10,
+ suffix: str = "",
+ save_image: bool = False,
+ max_size: int = -1,
+ ) -> Tuple:
+ """
+ Process a video for object segmentation and matting.
+ This method processes a video file by performing object segmentation and matting on each frame.
+ It supports warmup frames, mask erosion/dilation, and various output options.
+ Args:
+ input_path (str): Path to the input video file
+ mask_path (str): Path to the mask image file used for initial segmentation
+ output_path (str, optional): Directory path where output files will be saved. Defaults to a temporary directory
+ n_warmup (int, optional): Number of warmup frames to use. Defaults to 10
+ r_erode (int, optional): Erosion radius for mask processing. Defaults to 10
+ r_dilate (int, optional): Dilation radius for mask processing. Defaults to 10
+ suffix (str, optional): Suffix to append to output filename. Defaults to ""
+ save_image (bool, optional): Whether to save individual frames. Defaults to False
+ max_size (int, optional): Maximum size for frame dimension. Use -1 for no limit. Defaults to -1
+ Returns:
+ Tuple[str, str]: A tuple containing:
+ - Path to the output foreground video file (str)
+ - Path to the output alpha matte video file (str)
+ Output:
+ - Saves processed video files with foreground (_fgr) and alpha matte (_pha)
+ - If save_image=True, saves individual frames in separate directories
+ """
+ output_path = output_path if output_path is not None else tempfile.TemporaryDirectory().name
+ r_erode = int(r_erode)
+ r_dilate = int(r_dilate)
+ n_warmup = int(n_warmup)
+ max_size = int(max_size)
+
+ vframes, fps, width,height,length, video_name = read_frame_from_videos(input_path)
+ repeated_frames = vframes[0].unsqueeze(0).repeat(n_warmup, 1, 1, 1)
+ vframes = torch.cat([repeated_frames, vframes], dim=0).float()
+ length += n_warmup
+
+ new_h, new_w = vframes.shape[-2:]
+ if max_size > 0:
+ h, w = new_h, new_w
+ min_side = min(h, w)
+ if min_side > max_size:
+ new_h = int(h / min_side * max_size)
+ new_w = int(w / min_side * max_size)
+ vframes = F.interpolate(vframes, size=(new_h, new_w), mode="area")
+
+ os.makedirs(output_path, exist_ok=True)
+ if suffix:
+ video_name = f"{video_name}_{suffix}"
+ if save_image:
+ os.makedirs(f"{output_path}/{video_name}", exist_ok=True)
+ os.makedirs(f"{output_path}/{video_name}/pha", exist_ok=True)
+ os.makedirs(f"{output_path}/{video_name}/fgr", exist_ok=True)
+
+ # mask path 鍙互浼犲叆 mask image or np.array format object
+ if isinstance(mask_path,str):
+ mask = np.array(Image.open(mask_path).convert("L"))
+ else:
+ mask = mask_path
+
+ if r_dilate > 0:
+ mask = gen_dilate(mask, r_dilate, r_dilate)
+ if r_erode > 0:
+ mask = gen_erosion(mask, r_erode, r_erode)
+
+ mask = torch.from_numpy(mask).cuda()
+ if max_size > 0:
+ mask = F.interpolate(
+ mask.unsqueeze(0).unsqueeze(0), size=(new_h, new_w), mode="nearest"
+ )[0, 0]
+
+ bg_video_reader = None
+ if bg_video:
+ bg_video_reader = VideoReader(bg_video, transform=transforms.ToTensor())
+ elif bg_image:
+ bgr_IMAGE = Image.open(bg_image)
+ bgr = transforms.ToTensor()(bgr_IMAGE).to(self.device).unsqueeze(0).unsqueeze(0)
+ else:
+ # 浣跨敤 0,200,0 鑰屼笉鏄� 0,255,0 閬垮厤杩囧害鏇濆厜
+ bgr = torch.tensor([120, 255, 155], device=self.device, dtype=torch.float32).div(255).view(1, 1, 3, 1, 1)
+
+ #bgr = (np.array([120, 255, 155], dtype=np.float32) / 255).reshape((1, 1, 3))
+ objects = [1]
+
+ phas = []
+ fgrs = []
+
+ current_bg_idx = 0
+ bg_increasing = True
+ for ti in tqdm(range(length)):
+ #image = vframes[ti]
+ #image_np = np.array(image.permute(1, 2, 0))
+ #image = (image / 255.0).cuda().float()
+ image = vframes[ti].to(self.device).float() / 255.0 # GPU float32 [0,1]
+
+ if ti == 0:
+ output_prob = self.step(image, mask, objects=objects)
+ output_prob = self.step(image, first_frame_pred=True)
+ else:
+ if ti <= n_warmup:
+ output_prob = self.step(image, first_frame_pred=True)
+ else:
+ output_prob = self.step(image)
+
+ mask = self.output_prob_to_mask(output_prob)
+ #pha = mask.unsqueeze(2).cpu().numpy()
+ #com_np = image_np / 255.0 * pha + bgr * (1 - pha)
+ # GPU 铻嶅悎锛氭墿鎴� (1,1,H,W) 鈫� (1,1,1,H,W)
+ #pha_mask = mask.unsqueeze(0).unsqueeze(0).to(self.device) # [1,1,H,W]
+ #pha_mask = pha_mask.unsqueeze(2) # 鍙樻垚 (1, 1, 1, H, W)
+ #pha_mask = pha_mask.expand(-1, -1, bgr.shape[2], -1, -1) # 鍙樻垚 (1, 1, C, H, W)
+ pha_mask = mask[None, None, None].to(self.device).expand(-1, -1, 3, -1, -1)
+
+
+ # 鍓嶆櫙鍚堟垚锛堝叏閮� GPU锛�
+ if bg_video:
+ tmp_bgr = bg_video_reader[current_bg_idx].to(self.device)
+ if bg_increasing:
+ current_bg_idx += 1
+ if current_bg_idx >= len(bg_video_reader):
+ current_bg_idx = len(bg_video_reader) - 2
+ bg_increasing = False
+ else:
+ current_bg_idx -= 1
+ if current_bg_idx < 1:
+ current_bg_idx = 1
+ bg_increasing = True
+ com_t = image.unsqueeze(0).unsqueeze(0) * pha_mask + tmp_bgr * (1 - pha_mask) # GPU
+ else:
+ com_t = image.unsqueeze(0).unsqueeze(0) * pha_mask + bgr * (1 - pha_mask) # GPU
+
+ if ti > (n_warmup - 1):
+ #com_np = (com_np * 255).astype(np.uint8)
+ #pha = (pha * 255).astype(np.uint8)
+ # 淇濆瓨鍓嶅厛杞洖 CPU / NumPy
+ com_np = (com_t.squeeze().permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
+ pha_np = (pha_mask[0, 0, 0].cpu().numpy() * 255).astype(np.uint8)
+ pha_np = np.repeat(pha_np[:, :, None], 3, axis=2) # 鎻愬墠杞垚RGB鏍煎紡
+
+ fgrs.append(com_np)
+ phas.append(pha_np)
+ if save_image:
+ cv2.imwrite(
+ f"{output_path}/{video_name}/pha/{str(ti - n_warmup).zfill(5)}.png",
+ pha_np,
+ )
+ cv2.imwrite(
+ f"{output_path}/{video_name}/fgr/{str(ti - n_warmup).zfill(5)}.png",
+ com_np[..., [2, 1, 0]],
+ )
+
+ fgrs = np.array(fgrs)
+ phas = np.array(phas)
+
+ fgr_filename = f"{output_path}/{video_name}_fgr.mp4"
+ alpha_filename = f"{output_path}/{video_name}_pha.mp4"
+
+ writer_fgrs = VideoWriter(fgr_filename, fps, width,height,10_000_000,input_path)
+ writer_phas = VideoWriter(alpha_filename, fps, width,height,10_000_000,input_path)
+
+ for frame in fgrs:
+ writer_fgrs.write_nparray(frame)
+ for frame in phas:
+ writer_phas.write_nparray(frame)
+
+
+ #imageio.mimwrite(fgr_filename, fgrs, fps=fps, quality=7)
+ #imageio.mimwrite(alpha_filename, phas, fps=fps, quality=7)
+
+ writer_fgrs.close()
+ writer_phas.close()
+
+ return (fgr_filename,alpha_filename)
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/inference/kv_memory_store.py b/auto_media_publisher/modifiers/video_matting/matanyone/inference/kv_memory_store.py
new file mode 100644
index 0000000..e50b794
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/inference/kv_memory_store.py
@@ -0,0 +1,348 @@
+from typing import Dict, List, Optional, Literal
+from collections import defaultdict
+import torch
+
+
+def _add_last_dim(dictionary, key, new_value, prepend=False):
+ # append/prepend a new value to the last dimension of a tensor in a dictionary
+ # if the key does not exist, put the new value in
+ # append by default
+ if key in dictionary:
+ dictionary[key] = torch.cat([dictionary[key], new_value], -1)
+ else:
+ dictionary[key] = new_value
+
+
+class KeyValueMemoryStore:
+ """
+ Works for key/value pairs type storage
+ e.g., working and long-term memory
+ """
+ def __init__(self, save_selection: bool = False, save_usage: bool = False):
+ """
+ We store keys and values of objects that first appear in the same frame in a bucket.
+ Each bucket contains a set of object ids.
+ Each bucket is associated with a single key tensor
+ and a dictionary of value tensors indexed by object id.
+
+ The keys and values are stored as the concatenation of a permanent part and a temporary part.
+ """
+ self.save_selection = save_selection
+ self.save_usage = save_usage
+
+ self.global_bucket_id = 0 # does not reduce even if buckets are removed
+ self.buckets: Dict[int, List[int]] = {} # indexed by bucket id
+ self.k: Dict[int, torch.Tensor] = {} # indexed by bucket id
+ self.v: Dict[int, torch.Tensor] = {} # indexed by object id
+
+ # indexed by bucket id; the end point of permanent memory
+ self.perm_end_pt: Dict[int, int] = defaultdict(int)
+
+ # shrinkage and selection are just like the keys
+ self.s = {}
+ if self.save_selection:
+ self.e = {} # does not contain the permanent memory part
+
+ # usage
+ if self.save_usage:
+ self.use_cnt = {} # indexed by bucket id, does not contain the permanent memory part
+ self.life_cnt = {} # indexed by bucket id, does not contain the permanent memory part
+
+ def add(self,
+ key: torch.Tensor,
+ values: Dict[int, torch.Tensor],
+ shrinkage: torch.Tensor,
+ selection: torch.Tensor,
+ supposed_bucket_id: int = -1,
+ as_permanent: Literal['no', 'first', 'all'] = 'no') -> None:
+ """
+ key: (1/2)*C*N
+ values: dict of values ((1/2)*C*N), object ids are used as keys
+ shrinkage: (1/2)*1*N
+ selection: (1/2)*C*N
+
+ supposed_bucket_id: used to sync the bucket id between working and long-term memory
+ if provided, the input should all be in a single bucket indexed by this id
+ as_permanent: whether to store the input as permanent memory
+ 'no': don't
+ 'first': only store it as permanent memory if the bucket is empty
+ 'all': always store it as permanent memory
+ """
+ bs = key.shape[0]
+ ne = key.shape[-1]
+ assert len(key.shape) == 3
+ assert len(shrinkage.shape) == 3
+ assert not self.save_selection or len(selection.shape) == 3
+ assert as_permanent in ['no', 'first', 'all']
+
+ # add the value and create new buckets if necessary
+ if supposed_bucket_id >= 0:
+ enabled_buckets = [supposed_bucket_id]
+ bucket_exist = supposed_bucket_id in self.buckets
+ for obj, value in values.items():
+ if bucket_exist:
+ assert obj in self.v
+ assert obj in self.buckets[supposed_bucket_id]
+ _add_last_dim(self.v, obj, value, prepend=(as_permanent == 'all'))
+ else:
+ assert obj not in self.v
+ self.v[obj] = value
+ self.buckets[supposed_bucket_id] = list(values.keys())
+ else:
+ new_bucket_id = None
+ enabled_buckets = set()
+ for obj, value in values.items():
+ assert len(value.shape) == 3
+ if obj in self.v:
+ _add_last_dim(self.v, obj, value, prepend=(as_permanent == 'all'))
+ bucket_used = [
+ bucket_id for bucket_id, object_ids in self.buckets.items()
+ if obj in object_ids
+ ]
+ assert len(bucket_used) == 1 # each object should only be in one bucket
+ enabled_buckets.add(bucket_used[0])
+ else:
+ self.v[obj] = value
+ if new_bucket_id is None:
+ # create new bucket
+ new_bucket_id = self.global_bucket_id
+ self.global_bucket_id += 1
+ self.buckets[new_bucket_id] = []
+ # put the new object into the corresponding bucket
+ self.buckets[new_bucket_id].append(obj)
+ enabled_buckets.add(new_bucket_id)
+
+ # increment the permanent size if necessary
+ add_as_permanent = {} # indexed by bucket id
+ for bucket_id in enabled_buckets:
+ add_as_permanent[bucket_id] = False
+ if as_permanent == 'all':
+ self.perm_end_pt[bucket_id] += ne
+ add_as_permanent[bucket_id] = True
+ elif as_permanent == 'first':
+ if self.perm_end_pt[bucket_id] == 0:
+ self.perm_end_pt[bucket_id] = ne
+ add_as_permanent[bucket_id] = True
+
+ # create new counters for usage if necessary
+ if self.save_usage and as_permanent != 'all':
+ new_count = torch.zeros((bs, ne), device=key.device, dtype=torch.float32)
+ new_life = torch.zeros((bs, ne), device=key.device, dtype=torch.float32) + 1e-7
+
+ # add the key to every bucket
+ for bucket_id in self.buckets:
+ if bucket_id not in enabled_buckets:
+ # if we are not adding new values to a bucket, we should skip it
+ continue
+
+ _add_last_dim(self.k, bucket_id, key, prepend=add_as_permanent[bucket_id])
+ _add_last_dim(self.s, bucket_id, shrinkage, prepend=add_as_permanent[bucket_id])
+ if not add_as_permanent[bucket_id]:
+ if self.save_selection:
+ _add_last_dim(self.e, bucket_id, selection)
+ if self.save_usage:
+ _add_last_dim(self.use_cnt, bucket_id, new_count)
+ _add_last_dim(self.life_cnt, bucket_id, new_life)
+
+ def update_bucket_usage(self, bucket_id: int, usage: torch.Tensor) -> None:
+ # increase all life count by 1
+ # increase use of indexed elements
+ if not self.save_usage:
+ return
+
+ usage = usage[:, self.perm_end_pt[bucket_id]:]
+ if usage.shape[-1] == 0:
+ # if there is no temporary memory, we don't need to update
+ return
+ self.use_cnt[bucket_id] += usage.view_as(self.use_cnt[bucket_id])
+ self.life_cnt[bucket_id] += 1
+
+ def sieve_by_range(self, bucket_id: int, start: int, end: int, min_size: int) -> None:
+ # keep only the temporary elements *outside* of this range (with some boundary conditions)
+ # the permanent elements are ignored in this computation
+ # i.e., concat (a[:start], a[end:])
+ # bucket with size <= min_size are not modified
+
+ assert start >= 0
+ assert end <= 0
+
+ object_ids = self.buckets[bucket_id]
+ bucket_num_elements = self.k[bucket_id].shape[-1] - self.perm_end_pt[bucket_id]
+ if bucket_num_elements <= min_size:
+ return
+
+ if end == 0:
+ # negative 0 would not work as the end index!
+ # effectively make the second part an empty slice
+ end = self.k[bucket_id].shape[-1] + 1
+
+ p_size = self.perm_end_pt[bucket_id]
+ start = start + p_size
+
+ k = self.k[bucket_id]
+ s = self.s[bucket_id]
+ if self.save_selection:
+ e = self.e[bucket_id]
+ if self.save_usage:
+ use_cnt = self.use_cnt[bucket_id]
+ life_cnt = self.life_cnt[bucket_id]
+
+ self.k[bucket_id] = torch.cat([k[:, :, :start], k[:, :, end:]], -1)
+ self.s[bucket_id] = torch.cat([s[:, :, :start], s[:, :, end:]], -1)
+ if self.save_selection:
+ self.e[bucket_id] = torch.cat([e[:, :, :start - p_size], e[:, :, end:]], -1)
+ if self.save_usage:
+ self.use_cnt[bucket_id] = torch.cat([use_cnt[:, :start - p_size], use_cnt[:, end:]], -1)
+ self.life_cnt[bucket_id] = torch.cat([life_cnt[:, :start - p_size], life_cnt[:, end:]],
+ -1)
+ for obj_id in object_ids:
+ v = self.v[obj_id]
+ self.v[obj_id] = torch.cat([v[:, :, :start], v[:, :, end:]], -1)
+
+ def remove_old_memory(self, bucket_id: int, max_len: int) -> None:
+ self.sieve_by_range(bucket_id, 0, -max_len, max_len)
+
+ def remove_obsolete_features(self, bucket_id: int, max_size: int) -> None:
+ # for long-term memory only
+ object_ids = self.buckets[bucket_id]
+
+ assert self.perm_end_pt[bucket_id] == 0 # permanent memory should be empty in LT memory
+
+ # normalize with life duration
+ usage = self.get_usage(bucket_id)
+ bs = usage.shape[0]
+
+ survivals = []
+
+ for bi in range(bs):
+ _, survived = torch.topk(usage[bi], k=max_size)
+ survivals.append(survived.flatten())
+ assert survived.shape[-1] == survivals[0].shape[-1]
+
+ self.k[bucket_id] = torch.stack(
+ [self.k[bucket_id][bi, :, survived] for bi, survived in enumerate(survivals)], 0)
+ self.s[bucket_id] = torch.stack(
+ [self.s[bucket_id][bi, :, survived] for bi, survived in enumerate(survivals)], 0)
+
+ if self.save_selection:
+ # Long-term memory does not store selection so this should not be needed
+ self.e[bucket_id] = torch.stack(
+ [self.e[bucket_id][bi, :, survived] for bi, survived in enumerate(survivals)], 0)
+ for obj_id in object_ids:
+ self.v[obj_id] = torch.stack(
+ [self.v[obj_id][bi, :, survived] for bi, survived in enumerate(survivals)], 0)
+
+ self.use_cnt[bucket_id] = torch.stack(
+ [self.use_cnt[bucket_id][bi, survived] for bi, survived in enumerate(survivals)], 0)
+ self.life_cnt[bucket_id] = torch.stack(
+ [self.life_cnt[bucket_id][bi, survived] for bi, survived in enumerate(survivals)], 0)
+
+ def get_usage(self, bucket_id: int) -> torch.Tensor:
+ # return normalized usage
+ if not self.save_usage:
+ raise RuntimeError('I did not count usage!')
+ else:
+ usage = self.use_cnt[bucket_id] / self.life_cnt[bucket_id]
+ return usage
+
+ def get_all_sliced(
+ self, bucket_id: int, start: int, end: int
+ ) -> (torch.Tensor, torch.Tensor, torch.Tensor, Dict[int, torch.Tensor], torch.Tensor):
+ # return k, sk, ek, value, normalized usage in order, sliced by start and end
+ # this only queries the temporary memory
+
+ assert start >= 0
+ assert end <= 0
+
+ p_size = self.perm_end_pt[bucket_id]
+ start = start + p_size
+
+ if end == 0:
+ # negative 0 would not work as the end index!
+ k = self.k[bucket_id][:, :, start:]
+ sk = self.s[bucket_id][:, :, start:]
+ ek = self.e[bucket_id][:, :, start - p_size:] if self.save_selection else None
+ value = {obj_id: self.v[obj_id][:, :, start:] for obj_id in self.buckets[bucket_id]}
+ usage = self.get_usage(bucket_id)[:, start - p_size:] if self.save_usage else None
+ else:
+ k = self.k[bucket_id][:, :, start:end]
+ sk = self.s[bucket_id][:, :, start:end]
+ ek = self.e[bucket_id][:, :, start - p_size:end] if self.save_selection else None
+ value = {obj_id: self.v[obj_id][:, :, start:end] for obj_id in self.buckets[bucket_id]}
+ usage = self.get_usage(bucket_id)[:, start - p_size:end] if self.save_usage else None
+
+ return k, sk, ek, value, usage
+
+ def purge_except(self, obj_keep_idx: List[int]):
+ # purge certain objects from the memory except the one listed
+ obj_keep_idx = set(obj_keep_idx)
+
+ # remove objects that are not in the keep list from the buckets
+ buckets_to_remove = []
+ for bucket_id, object_ids in self.buckets.items():
+ self.buckets[bucket_id] = [obj_id for obj_id in object_ids if obj_id in obj_keep_idx]
+ if len(self.buckets[bucket_id]) == 0:
+ buckets_to_remove.append(bucket_id)
+
+ # remove object values that are not in the keep list
+ self.v = {k: v for k, v in self.v.items() if k in obj_keep_idx}
+
+ # remove buckets that are empty
+ for bucket_id in buckets_to_remove:
+ del self.buckets[bucket_id]
+ del self.k[bucket_id]
+ del self.s[bucket_id]
+ if self.save_selection:
+ del self.e[bucket_id]
+ if self.save_usage:
+ del self.use_cnt[bucket_id]
+ del self.life_cnt[bucket_id]
+
+ def clear_non_permanent_memory(self):
+ # clear all non-permanent memory
+ for bucket_id in self.buckets:
+ self.sieve_by_range(bucket_id, 0, 0, 0)
+
+ def get_v_size(self, obj_id: int) -> int:
+ return self.v[obj_id].shape[-1]
+
+ def size(self, bucket_id: int) -> int:
+ if bucket_id not in self.k:
+ return 0
+ else:
+ return self.k[bucket_id].shape[-1]
+
+ def perm_size(self, bucket_id: int) -> int:
+ return self.perm_end_pt[bucket_id]
+
+ def non_perm_size(self, bucket_id: int) -> int:
+ return self.size(bucket_id) - self.perm_size(bucket_id)
+
+ def engaged(self, bucket_id: Optional[int] = None) -> bool:
+ if bucket_id is None:
+ return len(self.buckets) > 0
+ else:
+ return bucket_id in self.buckets
+
+ @property
+ def num_objects(self) -> int:
+ return len(self.v)
+
+ @property
+ def key(self) -> Dict[int, torch.Tensor]:
+ return self.k
+
+ @property
+ def value(self) -> Dict[int, torch.Tensor]:
+ return self.v
+
+ @property
+ def shrinkage(self) -> Dict[int, torch.Tensor]:
+ return self.s
+
+ @property
+ def selection(self) -> Dict[int, torch.Tensor]:
+ return self.e
+
+ def __contains__(self, key):
+ return key in self.v
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/inference/memory_manager.py b/auto_media_publisher/modifiers/video_matting/matanyone/inference/memory_manager.py
new file mode 100644
index 0000000..1c8d9e1
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/inference/memory_manager.py
@@ -0,0 +1,453 @@
+import logging
+from omegaconf import DictConfig
+from typing import List, Dict
+import torch
+
+from auto_media_publisher.modifiers.video_matting.matanyone.inference.object_manager import ObjectManager
+from auto_media_publisher.modifiers.video_matting.matanyone.inference.kv_memory_store import KeyValueMemoryStore
+from auto_media_publisher.modifiers.video_matting.matanyone.model.matanyone import MatAnyone
+from auto_media_publisher.modifiers.video_matting.matanyone.model.utils.memory_utils import get_similarity, do_softmax
+
+log = logging.getLogger()
+
+
+class MemoryManager:
+ """
+ Manages all three memory stores and the transition between working/long-term memory
+ """
+ def __init__(self, cfg: DictConfig, object_manager: ObjectManager):
+ self.object_manager = object_manager
+ self.sensory_dim = cfg.model.sensory_dim
+ self.top_k = cfg.top_k
+ self.chunk_size = cfg.chunk_size
+
+ self.save_aux = cfg.save_aux
+
+ self.use_long_term = cfg.use_long_term
+ self.count_long_term_usage = cfg.long_term.count_usage
+ # subtract 1 because the first-frame is now counted as "permanent memory"
+ # and is not counted towards max_mem_frames
+ # but we want to keep the hyperparameters consistent as before for the same behavior
+ if self.use_long_term:
+ self.max_mem_frames = cfg.long_term.max_mem_frames - 1
+ self.min_mem_frames = cfg.long_term.min_mem_frames - 1
+ self.num_prototypes = cfg.long_term.num_prototypes
+ self.max_long_tokens = cfg.long_term.max_num_tokens
+ self.buffer_tokens = cfg.long_term.buffer_tokens
+ else:
+ self.max_mem_frames = cfg.max_mem_frames - 1
+
+ # dimensions will be inferred from input later
+ self.CK = self.CV = None
+ self.H = self.W = None
+
+ # The sensory memory is stored as a dictionary indexed by object ids
+ # each of shape bs * C^h * H * W
+ self.sensory = {}
+
+ # a dictionary indexed by object ids, each of shape bs * T * Q * C
+ self.obj_v = {}
+
+ self.work_mem = KeyValueMemoryStore(save_selection=self.use_long_term,
+ save_usage=self.use_long_term)
+ if self.use_long_term:
+ self.long_mem = KeyValueMemoryStore(save_usage=self.count_long_term_usage)
+
+ self.config_stale = True
+ self.engaged = False
+
+ def update_config(self, cfg: DictConfig) -> None:
+ self.config_stale = True
+ self.top_k = cfg['top_k']
+
+ assert self.use_long_term == cfg.use_long_term, 'cannot update this'
+ assert self.count_long_term_usage == cfg.long_term.count_usage, 'cannot update this'
+
+ self.use_long_term = cfg.use_long_term
+ self.count_long_term_usage = cfg.long_term.count_usage
+ if self.use_long_term:
+ self.max_mem_frames = cfg.long_term.max_mem_frames - 1
+ self.min_mem_frames = cfg.long_term.min_mem_frames - 1
+ self.num_prototypes = cfg.long_term.num_prototypes
+ self.max_long_tokens = cfg.long_term.max_num_tokens
+ self.buffer_tokens = cfg.long_term.buffer_tokens
+ else:
+ self.max_mem_frames = cfg.max_mem_frames - 1
+
+ def _readout(self, affinity, v, uncert_mask=None) -> torch.Tensor:
+ # affinity: bs*N*HW
+ # v: bs*C*N or bs*num_objects*C*N
+ # returns bs*C*HW or bs*num_objects*C*HW
+ if len(v.shape) == 3:
+ # single object
+ if uncert_mask is not None:
+ return v @ affinity * uncert_mask
+ else:
+ return v @ affinity
+ else:
+ bs, num_objects, C, N = v.shape
+ v = v.view(bs, num_objects * C, N)
+ out = v @ affinity
+ if uncert_mask is not None:
+ uncert_mask = uncert_mask.flatten(start_dim=2).expand(-1, C, -1)
+ out = out * uncert_mask
+ return out.view(bs, num_objects, C, -1)
+
+ def _get_mask_by_ids(self, mask: torch.Tensor, obj_ids: List[int]) -> torch.Tensor:
+ # -1 because the mask does not contain the background channel
+ return mask[:, [self.object_manager.find_tmp_by_id(obj) - 1 for obj in obj_ids]]
+
+ def _get_sensory_by_ids(self, obj_ids: List[int]) -> torch.Tensor:
+ return torch.stack([self.sensory[obj] for obj in obj_ids], dim=1)
+
+ def _get_object_mem_by_ids(self, obj_ids: List[int]) -> torch.Tensor:
+ return torch.stack([self.obj_v[obj] for obj in obj_ids], dim=1)
+
+ def _get_visual_values_by_ids(self, obj_ids: List[int]) -> torch.Tensor:
+ # All the values that the object ids refer to should have the same shape
+ value = torch.stack([self.work_mem.value[obj] for obj in obj_ids], dim=1)
+ if self.use_long_term and obj_ids[0] in self.long_mem.value:
+ lt_value = torch.stack([self.long_mem.value[obj] for obj in obj_ids], dim=1)
+ value = torch.cat([lt_value, value], dim=-1)
+
+ return value
+
+ def read_first_frame(self, last_msk_value, pix_feat: torch.Tensor,
+ last_mask: torch.Tensor, network: MatAnyone, uncert_output=None) -> Dict[int, torch.Tensor]:
+ """
+ Read from all memory stores and returns a single memory readout tensor for each object
+
+ pix_feat: (1/2) x C x H x W
+ query_key: (1/2) x C^k x H x W
+ selection: (1/2) x C^k x H x W
+ last_mask: (1/2) x num_objects x H x W (at stride 16)
+ return a dict of memory readouts, indexed by object indices. Each readout is C*H*W
+ """
+ h, w = pix_feat.shape[-2:]
+ bs = pix_feat.shape[0]
+ assert last_mask.shape[0] == bs
+
+ """
+ Compute affinity and perform readout
+ """
+ all_readout_mem = {}
+ buckets = self.work_mem.buckets
+ for bucket_id, bucket in buckets.items():
+
+ if self.chunk_size < 1:
+ object_chunks = [bucket]
+ else:
+ object_chunks = [
+ bucket[i:i + self.chunk_size] for i in range(0, len(bucket), self.chunk_size)
+ ]
+
+ for objects in object_chunks:
+ this_sensory = self._get_sensory_by_ids(objects)
+ this_last_mask = self._get_mask_by_ids(last_mask, objects)
+ this_msk_value = self._get_visual_values_by_ids(objects) # (1/2)*num_objects*C*N
+ pixel_readout = network.pixel_fusion(pix_feat, last_msk_value, this_sensory,
+ this_last_mask)
+ this_obj_mem = self._get_object_mem_by_ids(objects).unsqueeze(2)
+ readout_memory, aux_features = network.readout_query(pixel_readout, this_obj_mem)
+ for i, obj in enumerate(objects):
+ all_readout_mem[obj] = readout_memory[:, i]
+
+ if self.save_aux:
+ aux_output = {
+ # 'sensory': this_sensory,
+ # 'pixel_readout': pixel_readout,
+ 'q_logits': aux_features['logits'] if aux_features else None,
+ # 'q_weights': aux_features['q_weights'] if aux_features else None,
+ # 'p_weights': aux_features['p_weights'] if aux_features else None,
+ # 'attn_mask': aux_features['attn_mask'].float() if aux_features else None,
+ }
+ self.aux = aux_output
+
+ return all_readout_mem
+
+ def read(self, pix_feat: torch.Tensor, query_key: torch.Tensor, selection: torch.Tensor,
+ last_mask: torch.Tensor, network: MatAnyone, uncert_output=None, last_msk_value=None, ti=None,
+ last_pix_feat=None, last_pred_mask=None) -> Dict[int, torch.Tensor]:
+ """
+ Read from all memory stores and returns a single memory readout tensor for each object
+
+ pix_feat: (1/2) x C x H x W
+ query_key: (1/2) x C^k x H x W
+ selection: (1/2) x C^k x H x W
+ last_mask: (1/2) x num_objects x H x W (at stride 16)
+ return a dict of memory readouts, indexed by object indices. Each readout is C*H*W
+ """
+ h, w = pix_feat.shape[-2:]
+ bs = pix_feat.shape[0]
+ assert query_key.shape[0] == bs
+ assert selection.shape[0] == bs
+ assert last_mask.shape[0] == bs
+
+ uncert_mask = uncert_output["mask"] if uncert_output is not None else None
+
+ query_key = query_key.flatten(start_dim=2) # bs*C^k*HW
+ selection = selection.flatten(start_dim=2) # bs*C^k*HW
+ """
+ Compute affinity and perform readout
+ """
+ all_readout_mem = {}
+ buckets = self.work_mem.buckets
+ for bucket_id, bucket in buckets.items():
+ if self.use_long_term and self.long_mem.engaged(bucket_id):
+ # Use long-term memory
+ long_mem_size = self.long_mem.size(bucket_id)
+ memory_key = torch.cat([self.long_mem.key[bucket_id], self.work_mem.key[bucket_id]],
+ -1)
+ shrinkage = torch.cat(
+ [self.long_mem.shrinkage[bucket_id], self.work_mem.shrinkage[bucket_id]], -1)
+
+ similarity = get_similarity(memory_key, shrinkage, query_key, selection)
+ affinity, usage = do_softmax(similarity,
+ top_k=self.top_k,
+ inplace=True,
+ return_usage=True)
+ """
+ Record memory usage for working and long-term memory
+ """
+ # ignore the index return for long-term memory
+ work_usage = usage[:, long_mem_size:]
+ self.work_mem.update_bucket_usage(bucket_id, work_usage)
+
+ if self.count_long_term_usage:
+ # ignore the index return for working memory
+ long_usage = usage[:, :long_mem_size]
+ self.long_mem.update_bucket_usage(bucket_id, long_usage)
+ else:
+ # no long-term memory
+ memory_key = self.work_mem.key[bucket_id]
+ shrinkage = self.work_mem.shrinkage[bucket_id]
+ similarity = get_similarity(memory_key, shrinkage, query_key, selection, uncert_mask=uncert_mask)
+
+ if self.use_long_term:
+ affinity, usage = do_softmax(similarity,
+ top_k=self.top_k,
+ inplace=True,
+ return_usage=True)
+ self.work_mem.update_bucket_usage(bucket_id, usage)
+ else:
+ affinity = do_softmax(similarity, top_k=self.top_k, inplace=True)
+
+ if self.chunk_size < 1:
+ object_chunks = [bucket]
+ else:
+ object_chunks = [
+ bucket[i:i + self.chunk_size] for i in range(0, len(bucket), self.chunk_size)
+ ]
+
+ for objects in object_chunks:
+ this_sensory = self._get_sensory_by_ids(objects)
+ this_last_mask = self._get_mask_by_ids(last_mask, objects)
+ this_msk_value = self._get_visual_values_by_ids(objects) # (1/2)*num_objects*C*N
+ visual_readout = self._readout(affinity,
+ this_msk_value, uncert_mask).view(bs, len(objects), self.CV, h, w)
+
+ uncert_output = network.pred_uncertainty(last_pix_feat, pix_feat, last_pred_mask, visual_readout[:,0]-last_msk_value[:,0])
+
+ if uncert_output is not None:
+ uncert_prob = uncert_output["prob"].unsqueeze(1) # b n 1 h w
+ visual_readout = visual_readout*uncert_prob + last_msk_value*(1-uncert_prob)
+
+ pixel_readout = network.pixel_fusion(pix_feat, visual_readout, this_sensory,
+ this_last_mask)
+ this_obj_mem = self._get_object_mem_by_ids(objects).unsqueeze(2)
+ readout_memory, aux_features = network.readout_query(pixel_readout, this_obj_mem)
+ for i, obj in enumerate(objects):
+ all_readout_mem[obj] = readout_memory[:, i]
+
+ if self.save_aux:
+ aux_output = {
+ # 'sensory': this_sensory,
+ # 'pixel_readout': pixel_readout,
+ 'q_logits': aux_features['logits'] if aux_features else None,
+ # 'q_weights': aux_features['q_weights'] if aux_features else None,
+ # 'p_weights': aux_features['p_weights'] if aux_features else None,
+ # 'attn_mask': aux_features['attn_mask'].float() if aux_features else None,
+ }
+ self.aux = aux_output
+
+ return all_readout_mem
+
+ def add_memory(self,
+ key: torch.Tensor,
+ shrinkage: torch.Tensor,
+ msk_value: torch.Tensor,
+ obj_value: torch.Tensor,
+ objects: List[int],
+ selection: torch.Tensor = None,
+ *,
+ as_permanent: bool = False) -> None:
+ # key: (1/2)*C*H*W
+ # msk_value: (1/2)*num_objects*C*H*W
+ # obj_value: (1/2)*num_objects*Q*C
+ # objects contains a list of object ids corresponding to the objects in msk_value/obj_value
+ bs = key.shape[0]
+ assert shrinkage.shape[0] == bs
+ assert msk_value.shape[0] == bs
+ assert obj_value.shape[0] == bs
+
+ self.engaged = True
+ if self.H is None or self.config_stale:
+ self.config_stale = False
+ self.H, self.W = msk_value.shape[-2:]
+ self.HW = self.H * self.W
+ # convert from num. frames to num. tokens
+ self.max_work_tokens = self.max_mem_frames * self.HW
+ if self.use_long_term:
+ self.min_work_tokens = self.min_mem_frames * self.HW
+
+ # key: bs*C*N
+ # value: bs*num_objects*C*N
+ key = key.flatten(start_dim=2)
+ shrinkage = shrinkage.flatten(start_dim=2)
+ self.CK = key.shape[1]
+
+ msk_value = msk_value.flatten(start_dim=3)
+ self.CV = msk_value.shape[2]
+
+ if selection is not None:
+ # not used in non-long-term mode
+ selection = selection.flatten(start_dim=2)
+
+ # insert object values into object memory
+ for obj_id, obj in enumerate(objects):
+ if obj in self.obj_v:
+ """streaming average
+ each self.obj_v[obj] is (1/2)*num_summaries*(embed_dim+1)
+ first embed_dim keeps track of the sum of embeddings
+ the last dim keeps the total count
+ averaging in done inside the object transformer
+
+ incoming obj_value is (1/2)*num_objects*num_summaries*(embed_dim+1)
+ self.obj_v[obj] = torch.cat([self.obj_v[obj], obj_value[:, obj_id]], dim=0)
+ """
+ last_acc = self.obj_v[obj][:, :, -1]
+ new_acc = last_acc + obj_value[:, obj_id, :, -1]
+
+ self.obj_v[obj][:, :, :-1] = (self.obj_v[obj][:, :, :-1] +
+ obj_value[:, obj_id, :, :-1])
+ self.obj_v[obj][:, :, -1] = new_acc
+ else:
+ self.obj_v[obj] = obj_value[:, obj_id]
+
+ # convert mask value tensor into a dict for insertion
+ msk_values = {obj: msk_value[:, obj_id] for obj_id, obj in enumerate(objects)}
+ self.work_mem.add(key,
+ msk_values,
+ shrinkage,
+ selection=selection,
+ as_permanent=as_permanent)
+
+ for bucket_id in self.work_mem.buckets.keys():
+ # long-term memory cleanup
+ if self.use_long_term:
+ # Do memory compressed if needed
+ if self.work_mem.non_perm_size(bucket_id) >= self.max_work_tokens:
+ # Remove obsolete features if needed
+ if self.long_mem.non_perm_size(bucket_id) >= (self.max_long_tokens -
+ self.num_prototypes):
+ self.long_mem.remove_obsolete_features(
+ bucket_id,
+ self.max_long_tokens - self.num_prototypes - self.buffer_tokens)
+
+ self.compress_features(bucket_id)
+ else:
+ # FIFO
+ self.work_mem.remove_old_memory(bucket_id, self.max_work_tokens)
+
+ def purge_except(self, obj_keep_idx: List[int]) -> None:
+ # purge certain objects from the memory except the one listed
+ self.work_mem.purge_except(obj_keep_idx)
+ if self.use_long_term and self.long_mem.engaged():
+ self.long_mem.purge_except(obj_keep_idx)
+ self.sensory = {k: v for k, v in self.sensory.items() if k in obj_keep_idx}
+
+ if not self.work_mem.engaged():
+ # everything is removed!
+ self.engaged = False
+
+ def compress_features(self, bucket_id: int) -> None:
+
+ # perform memory consolidation
+ prototype_key, prototype_value, prototype_shrinkage = self.consolidation(
+ *self.work_mem.get_all_sliced(bucket_id, 0, -self.min_work_tokens))
+
+ # remove consolidated working memory
+ self.work_mem.sieve_by_range(bucket_id,
+ 0,
+ -self.min_work_tokens,
+ min_size=self.min_work_tokens)
+
+ # add to long-term memory
+ self.long_mem.add(prototype_key,
+ prototype_value,
+ prototype_shrinkage,
+ selection=None,
+ supposed_bucket_id=bucket_id)
+
+ def consolidation(self, candidate_key: torch.Tensor, candidate_shrinkage: torch.Tensor,
+ candidate_selection: torch.Tensor, candidate_value: Dict[int, torch.Tensor],
+ usage: torch.Tensor) -> (torch.Tensor, Dict[int, torch.Tensor], torch.Tensor):
+ # find the indices with max usage
+ bs = candidate_key.shape[0]
+ assert bs in [1, 2]
+
+ prototype_key = []
+ prototype_selection = []
+ for bi in range(bs):
+ _, max_usage_indices = torch.topk(usage[bi], k=self.num_prototypes, dim=-1, sorted=True)
+ prototype_indices = max_usage_indices.flatten()
+ prototype_key.append(candidate_key[bi, :, prototype_indices])
+ prototype_selection.append(candidate_selection[bi, :, prototype_indices])
+ prototype_key = torch.stack(prototype_key, dim=0)
+ prototype_selection = torch.stack(prototype_selection, dim=0)
+ """
+ Potentiation step
+ """
+ similarity = get_similarity(candidate_key, candidate_shrinkage, prototype_key,
+ prototype_selection)
+ affinity = do_softmax(similarity)
+
+ # readout the values
+ prototype_value = {k: self._readout(affinity, v) for k, v in candidate_value.items()}
+
+ # readout the shrinkage term
+ prototype_shrinkage = self._readout(affinity, candidate_shrinkage)
+
+ return prototype_key, prototype_value, prototype_shrinkage
+
+ def initialize_sensory_if_needed(self, sample_key: torch.Tensor, ids: List[int]):
+ for obj in ids:
+ if obj not in self.sensory:
+ # also initializes the sensory memory
+ bs, _, h, w = sample_key.shape
+ self.sensory[obj] = torch.zeros((bs, self.sensory_dim, h, w),
+ device=sample_key.device)
+
+ def update_sensory(self, sensory: torch.Tensor, ids: List[int]):
+ # sensory: 1*num_objects*C*H*W
+ for obj_id, obj in enumerate(ids):
+ self.sensory[obj] = sensory[:, obj_id]
+
+ def get_sensory(self, ids: List[int]):
+ # returns (1/2)*num_objects*C*H*W
+ return self._get_sensory_by_ids(ids)
+
+ def clear_non_permanent_memory(self):
+ self.work_mem.clear_non_permanent_memory()
+ if self.use_long_term:
+ self.long_mem.clear_non_permanent_memory()
+
+ def clear_sensory_memory(self):
+ self.sensory = {}
+
+ def clear_work_mem(self):
+ self.work_mem = KeyValueMemoryStore(save_selection=self.use_long_term,
+ save_usage=self.use_long_term)
+
+ def clear_obj_mem(self):
+ self.obj_v = {}
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/inference/object_info.py b/auto_media_publisher/modifiers/video_matting/matanyone/inference/object_info.py
new file mode 100644
index 0000000..b0e0bd4
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/inference/object_info.py
@@ -0,0 +1,24 @@
+class ObjectInfo:
+ """
+ Store meta information for an object
+ """
+ def __init__(self, id: int):
+ self.id = id
+ self.poke_count = 0 # count number of detections missed
+
+ def poke(self) -> None:
+ self.poke_count += 1
+
+ def unpoke(self) -> None:
+ self.poke_count = 0
+
+ def __hash__(self):
+ return hash(self.id)
+
+ def __eq__(self, other):
+ if type(other) == int:
+ return self.id == other
+ return self.id == other.id
+
+ def __repr__(self):
+ return f'(ID: {self.id})'
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/inference/object_manager.py b/auto_media_publisher/modifiers/video_matting/matanyone/inference/object_manager.py
new file mode 100644
index 0000000..30c6caa
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/inference/object_manager.py
@@ -0,0 +1,149 @@
+from typing import Union, List, Dict
+
+import torch
+from auto_media_publisher.modifiers.video_matting.matanyone.inference.object_info import ObjectInfo
+
+
+class ObjectManager:
+ """
+ Object IDs are immutable. The same ID always represent the same object.
+ Temporary IDs are the positions of each object in the tensor. It changes as objects get removed.
+ Temporary IDs start from 1.
+ """
+
+ def __init__(self):
+ self.obj_to_tmp_id: Dict[ObjectInfo, int] = {}
+ self.tmp_id_to_obj: Dict[int, ObjectInfo] = {}
+ self.obj_id_to_obj: Dict[int, ObjectInfo] = {}
+
+ self.all_historical_object_ids: List[int] = []
+
+ def _recompute_obj_id_to_obj_mapping(self) -> None:
+ self.obj_id_to_obj = {obj.id: obj for obj in self.obj_to_tmp_id}
+
+ def add_new_objects(
+ self, objects: Union[List[ObjectInfo], ObjectInfo,
+ List[int]]) -> (List[int], List[int]):
+ if not isinstance(objects, list):
+ objects = [objects]
+
+ corresponding_tmp_ids = []
+ corresponding_obj_ids = []
+ for obj in objects:
+ if isinstance(obj, int):
+ obj = ObjectInfo(id=obj)
+
+ if obj in self.obj_to_tmp_id:
+ # old object
+ corresponding_tmp_ids.append(self.obj_to_tmp_id[obj])
+ corresponding_obj_ids.append(obj.id)
+ else:
+ # new object
+ new_obj = ObjectInfo(id=obj.id)
+
+ # new object
+ new_tmp_id = len(self.obj_to_tmp_id) + 1
+ self.obj_to_tmp_id[new_obj] = new_tmp_id
+ self.tmp_id_to_obj[new_tmp_id] = new_obj
+ self.all_historical_object_ids.append(new_obj.id)
+ corresponding_tmp_ids.append(new_tmp_id)
+ corresponding_obj_ids.append(new_obj.id)
+
+ self._recompute_obj_id_to_obj_mapping()
+ assert corresponding_tmp_ids == sorted(corresponding_tmp_ids)
+ return corresponding_tmp_ids, corresponding_obj_ids
+
+ def delete_objects(self, obj_ids_to_remove: Union[int, List[int]]) -> None:
+ # delete an object or a list of objects
+ # re-sort the tmp ids
+ if isinstance(obj_ids_to_remove, int):
+ obj_ids_to_remove = [obj_ids_to_remove]
+
+ new_tmp_id = 1
+ total_num_id = len(self.obj_to_tmp_id)
+
+ local_obj_to_tmp_id = {}
+ local_tmp_to_obj_id = {}
+
+ for tmp_iter in range(1, total_num_id + 1):
+ obj = self.tmp_id_to_obj[tmp_iter]
+ if obj.id not in obj_ids_to_remove:
+ local_obj_to_tmp_id[obj] = new_tmp_id
+ local_tmp_to_obj_id[new_tmp_id] = obj
+ new_tmp_id += 1
+
+ self.obj_to_tmp_id = local_obj_to_tmp_id
+ self.tmp_id_to_obj = local_tmp_to_obj_id
+ self._recompute_obj_id_to_obj_mapping()
+
+ def purge_inactive_objects(self,
+ max_missed_detection_count: int) -> (bool, List[int], List[int]):
+ # remove tmp ids of objects that are removed
+ obj_id_to_be_deleted = []
+ tmp_id_to_be_deleted = []
+ tmp_id_to_keep = []
+ obj_id_to_keep = []
+
+ for obj in self.obj_to_tmp_id:
+ if obj.poke_count > max_missed_detection_count:
+ obj_id_to_be_deleted.append(obj.id)
+ tmp_id_to_be_deleted.append(self.obj_to_tmp_id[obj])
+ else:
+ tmp_id_to_keep.append(self.obj_to_tmp_id[obj])
+ obj_id_to_keep.append(obj.id)
+
+ purge_activated = len(obj_id_to_be_deleted) > 0
+ if purge_activated:
+ self.delete_objects(obj_id_to_be_deleted)
+ return purge_activated, tmp_id_to_keep, obj_id_to_keep
+
+ def tmp_to_obj_cls(self, mask) -> torch.Tensor:
+ # remap tmp id cls representation to the true object id representation
+ new_mask = torch.zeros_like(mask)
+ for tmp_id, obj in self.tmp_id_to_obj.items():
+ new_mask[mask == tmp_id] = obj.id
+ return new_mask
+
+ def get_tmp_to_obj_mapping(self) -> Dict[int, ObjectInfo]:
+ # returns the mapping in a dict format for saving it with pickle
+ return {obj.id: tmp_id for obj, tmp_id in self.tmp_id_to_obj.items()}
+
+ def realize_dict(self, obj_dict, dim=1) -> torch.Tensor:
+ # turns a dict indexed by obj id into a tensor, ordered by tmp IDs
+ output = []
+ for _, obj in self.tmp_id_to_obj.items():
+ if obj.id not in obj_dict:
+ raise NotImplementedError
+ output.append(obj_dict[obj.id])
+ output = torch.stack(output, dim=dim)
+ return output
+
+ def make_one_hot(self, cls_mask) -> torch.Tensor:
+ output = []
+ for _, obj in self.tmp_id_to_obj.items():
+ output.append(cls_mask == obj.id)
+ if len(output) == 0:
+ output = torch.zeros((0, *cls_mask.shape), dtype=torch.bool, device=cls_mask.device)
+ else:
+ output = torch.stack(output, dim=0)
+ return output
+
+ @property
+ def all_obj_ids(self) -> List[int]:
+ return [k.id for k in self.obj_to_tmp_id]
+
+ @property
+ def num_obj(self) -> int:
+ return len(self.obj_to_tmp_id)
+
+ def has_all(self, objects: List[int]) -> bool:
+ for obj in objects:
+ if obj not in self.obj_to_tmp_id:
+ return False
+ return True
+
+ def find_object_by_id(self, obj_id) -> ObjectInfo:
+ return self.obj_id_to_obj[obj_id]
+
+ def find_tmp_by_id(self, obj_id) -> int:
+ return self.obj_to_tmp_id[self.obj_id_to_obj[obj_id]]
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/__init__.py b/auto_media_publisher/modifiers/video_matting/matanyone/inference/utils/__init__.py
similarity index 100%
copy from auto_media_publisher/modifiers/video_matting/matting_model/__init__.py
copy to auto_media_publisher/modifiers/video_matting/matanyone/inference/utils/__init__.py
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/inference/utils/args_utils.py b/auto_media_publisher/modifiers/video_matting/matanyone/inference/utils/args_utils.py
new file mode 100644
index 0000000..a771cca
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/inference/utils/args_utils.py
@@ -0,0 +1,30 @@
+import logging
+from omegaconf import DictConfig
+
+log = logging.getLogger()
+
+
+def get_dataset_cfg(cfg: DictConfig):
+ dataset_name = cfg.dataset
+ data_cfg = cfg.datasets[dataset_name]
+
+ potential_overrides = [
+ 'image_directory',
+ 'mask_directory',
+ 'json_directory',
+ 'size',
+ 'save_all',
+ 'use_all_masks',
+ 'use_long_term',
+ 'mem_every',
+ ]
+
+ for override in potential_overrides:
+ if cfg[override] is not None:
+ log.info(f'Overriding config {override} from {data_cfg[override]} to {cfg[override]}')
+ data_cfg[override] = cfg[override]
+ # escalte all potential overrides to the top-level config
+ if override in data_cfg:
+ cfg[override] = data_cfg[override]
+
+ return data_cfg
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/__init__.py b/auto_media_publisher/modifiers/video_matting/matanyone/model/__init__.py
similarity index 100%
rename from auto_media_publisher/modifiers/video_matting/matting_model/__init__.py
rename to auto_media_publisher/modifiers/video_matting/matanyone/model/__init__.py
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/model/aux_modules.py b/auto_media_publisher/modifiers/video_matting/matanyone/model/aux_modules.py
new file mode 100644
index 0000000..2e9f627
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/model/aux_modules.py
@@ -0,0 +1,93 @@
+"""
+For computing auxiliary outputs for auxiliary losses
+"""
+from typing import Dict
+from omegaconf import DictConfig
+import torch
+import torch.nn as nn
+
+from auto_media_publisher.modifiers.video_matting.matanyone.model.group_modules import GConv2d
+from auto_media_publisher.modifiers.video_matting.matanyone.utils.tensor_utils import aggregate
+
+
+class LinearPredictor(nn.Module):
+ def __init__(self, x_dim: int, pix_dim: int):
+ super().__init__()
+ self.projection = GConv2d(x_dim, pix_dim + 1, kernel_size=1)
+
+ def forward(self, pix_feat: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+ # pixel_feat: B*pix_dim*H*W
+ # x: B*num_objects*x_dim*H*W
+ num_objects = x.shape[1]
+ x = self.projection(x)
+
+ pix_feat = pix_feat.unsqueeze(1).expand(-1, num_objects, -1, -1, -1)
+ logits = (pix_feat * x[:, :, :-1]).sum(dim=2) + x[:, :, -1]
+ return logits
+
+
+class DirectPredictor(nn.Module):
+ def __init__(self, x_dim: int):
+ super().__init__()
+ self.projection = GConv2d(x_dim, 1, kernel_size=1)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ # x: B*num_objects*x_dim*H*W
+ logits = self.projection(x).squeeze(2)
+ return logits
+
+
+class AuxComputer(nn.Module):
+ def __init__(self, cfg: DictConfig):
+ super().__init__()
+
+ use_sensory_aux = cfg.model.aux_loss.sensory.enabled
+ self.use_query_aux = cfg.model.aux_loss.query.enabled
+ self.use_sensory_aux = use_sensory_aux
+
+ sensory_dim = cfg.model.sensory_dim
+ embed_dim = cfg.model.embed_dim
+
+ if use_sensory_aux:
+ self.sensory_aux = LinearPredictor(sensory_dim, embed_dim)
+
+ def _aggregate_with_selector(self, logits: torch.Tensor, selector: torch.Tensor) -> torch.Tensor:
+ prob = torch.sigmoid(logits)
+ if selector is not None:
+ prob = prob * selector
+ logits = aggregate(prob, dim=1)
+ return logits
+
+ def forward(self, pix_feat: torch.Tensor, aux_input: Dict[str, torch.Tensor],
+ selector: torch.Tensor, seg_pass=False) -> Dict[str, torch.Tensor]:
+ sensory = aux_input['sensory']
+ q_logits = aux_input['q_logits']
+
+ aux_output = {}
+ aux_output['attn_mask'] = aux_input['attn_mask']
+
+ if self.use_sensory_aux:
+ # B*num_objects*H*W
+ logits = self.sensory_aux(pix_feat, sensory)
+ aux_output['sensory_logits'] = self._aggregate_with_selector(logits, selector)
+ if self.use_query_aux:
+ # B*num_objects*num_levels*H*W
+ aux_output['q_logits'] = self._aggregate_with_selector(
+ torch.stack(q_logits, dim=2),
+ selector.unsqueeze(2) if selector is not None else None)
+
+ return aux_output
+
+ def compute_mask(self, aux_input: Dict[str, torch.Tensor],
+ selector: torch.Tensor) -> Dict[str, torch.Tensor]:
+ # sensory = aux_input['sensory']
+ q_logits = aux_input['q_logits']
+
+ aux_output = {}
+
+ # B*num_objects*num_levels*H*W
+ aux_output['q_logits'] = self._aggregate_with_selector(
+ torch.stack(q_logits, dim=2),
+ selector.unsqueeze(2) if selector is not None else None)
+
+ return aux_output
\ No newline at end of file
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/model/big_modules.py b/auto_media_publisher/modifiers/video_matting/matanyone/model/big_modules.py
new file mode 100644
index 0000000..4dbf90d
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/model/big_modules.py
@@ -0,0 +1,365 @@
+"""
+big_modules.py - This file stores higher-level network blocks.
+
+x - usually denotes features that are shared between objects.
+g - usually denotes features that are not shared between objects
+ with an extra "num_objects" dimension (batch_size * num_objects * num_channels * H * W).
+
+The trailing number of a variable usually denotes the stride
+"""
+
+from typing import Iterable
+from omegaconf import DictConfig
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from auto_media_publisher.modifiers.video_matting.matanyone.model.group_modules import MainToGroupDistributor, GroupFeatureFusionBlock, GConv2d
+from auto_media_publisher.modifiers.video_matting.matanyone.model.utils import resnet
+from auto_media_publisher.modifiers.video_matting.matanyone.model.modules import SensoryDeepUpdater, SensoryUpdater_fullscale, DecoderFeatureProcessor, MaskUpsampleBlock
+
+class UncertPred(nn.Module):
+ def __init__(self, model_cfg: DictConfig):
+ super().__init__()
+ self.conv1x1_v2 = nn.Conv2d(model_cfg.pixel_dim*2 + 1 + model_cfg.value_dim, 64, kernel_size=1, stride=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(64)
+ self.relu = nn.ReLU(inplace=True)
+ self.conv3x3 = nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1, groups=1, bias=False, dilation=1)
+ self.bn2 = nn.BatchNorm2d(32)
+ self.conv3x3_out = nn.Conv2d(32, 1, kernel_size=3, stride=1, padding=1, groups=1, bias=False, dilation=1)
+
+ def forward(self, last_frame_feat: torch.Tensor, cur_frame_feat: torch.Tensor, last_mask: torch.Tensor, mem_val_diff:torch.Tensor):
+ last_mask = F.interpolate(last_mask, size=last_frame_feat.shape[-2:], mode='area')
+ x = torch.cat([last_frame_feat, cur_frame_feat, last_mask, mem_val_diff], dim=1)
+ x = self.conv1x1_v2(x)
+ x = self.bn1(x)
+ x = self.relu(x)
+ x = self.conv3x3(x)
+ x = self.bn2(x)
+ x = self.relu(x)
+ x = self.conv3x3_out(x)
+ return x
+
+ # override the default train() to freeze BN statistics
+ def train(self, mode=True):
+ self.training = False
+ for module in self.children():
+ module.train(False)
+ return self
+
+class PixelEncoder(nn.Module):
+ def __init__(self, model_cfg: DictConfig):
+ super().__init__()
+
+ self.is_resnet = 'resnet' in model_cfg.pixel_encoder.type
+ # if model_cfg.pretrained_resnet is set in the model_cfg we get the value
+ # else default to True
+ is_pretrained_resnet = getattr(model_cfg,"pretrained_resnet",True)
+ if self.is_resnet:
+ if model_cfg.pixel_encoder.type == 'resnet18':
+ network = resnet.resnet18(pretrained=is_pretrained_resnet)
+ elif model_cfg.pixel_encoder.type == 'resnet50':
+ network = resnet.resnet50(pretrained=is_pretrained_resnet)
+ else:
+ raise NotImplementedError
+ self.conv1 = network.conv1
+ self.bn1 = network.bn1
+ self.relu = network.relu
+ self.maxpool = network.maxpool
+
+ self.res2 = network.layer1
+ self.layer2 = network.layer2
+ self.layer3 = network.layer3
+ else:
+ raise NotImplementedError
+
+ def forward(self, x: torch.Tensor, seq_length=None) -> (torch.Tensor, torch.Tensor, torch.Tensor):
+ f1 = x
+ x = self.conv1(x)
+ x = self.bn1(x)
+ x = self.relu(x)
+ f2 = x
+ x = self.maxpool(x)
+ f4 = self.res2(x)
+ f8 = self.layer2(f4)
+ f16 = self.layer3(f8)
+
+ return f16, f8, f4, f2, f1
+
+ # override the default train() to freeze BN statistics
+ def train(self, mode=True):
+ self.training = False
+ for module in self.children():
+ module.train(False)
+ return self
+
+
+class KeyProjection(nn.Module):
+ def __init__(self, model_cfg: DictConfig):
+ super().__init__()
+ in_dim = model_cfg.pixel_encoder.ms_dims[0]
+ mid_dim = model_cfg.pixel_dim
+ key_dim = model_cfg.key_dim
+
+ self.pix_feat_proj = nn.Conv2d(in_dim, mid_dim, kernel_size=1)
+ self.key_proj = nn.Conv2d(mid_dim, key_dim, kernel_size=3, padding=1)
+ # shrinkage
+ self.d_proj = nn.Conv2d(mid_dim, 1, kernel_size=3, padding=1)
+ # selection
+ self.e_proj = nn.Conv2d(mid_dim, key_dim, kernel_size=3, padding=1)
+
+ nn.init.orthogonal_(self.key_proj.weight.data)
+ nn.init.zeros_(self.key_proj.bias.data)
+
+ def forward(self, x: torch.Tensor, *, need_s: bool,
+ need_e: bool) -> (torch.Tensor, torch.Tensor, torch.Tensor):
+ x = self.pix_feat_proj(x)
+ shrinkage = self.d_proj(x)**2 + 1 if (need_s) else None
+ selection = torch.sigmoid(self.e_proj(x)) if (need_e) else None
+
+ return self.key_proj(x), shrinkage, selection
+
+
+class MaskEncoder(nn.Module):
+ def __init__(self, model_cfg: DictConfig, single_object=False):
+ super().__init__()
+ pixel_dim = model_cfg.pixel_dim
+ value_dim = model_cfg.value_dim
+ sensory_dim = model_cfg.sensory_dim
+ final_dim = model_cfg.mask_encoder.final_dim
+
+ self.single_object = single_object
+ extra_dim = 1 if single_object else 2
+
+ # if model_cfg.pretrained_resnet is set in the model_cfg we get the value
+ # else default to True
+ is_pretrained_resnet = getattr(model_cfg,"pretrained_resnet",True)
+ if model_cfg.mask_encoder.type == 'resnet18':
+ network = resnet.resnet18(pretrained=is_pretrained_resnet, extra_dim=extra_dim)
+ elif model_cfg.mask_encoder.type == 'resnet50':
+ network = resnet.resnet50(pretrained=is_pretrained_resnet, extra_dim=extra_dim)
+ else:
+ raise NotImplementedError
+ self.conv1 = network.conv1
+ self.bn1 = network.bn1
+ self.relu = network.relu
+ self.maxpool = network.maxpool
+
+ self.layer1 = network.layer1
+ self.layer2 = network.layer2
+ self.layer3 = network.layer3
+
+ self.distributor = MainToGroupDistributor()
+ self.fuser = GroupFeatureFusionBlock(pixel_dim, final_dim, value_dim)
+
+ self.sensory_update = SensoryDeepUpdater(value_dim, sensory_dim)
+
+ def forward(self,
+ image: torch.Tensor,
+ pix_feat: torch.Tensor,
+ sensory: torch.Tensor,
+ masks: torch.Tensor,
+ others: torch.Tensor,
+ *,
+ deep_update: bool = True,
+ chunk_size: int = -1) -> (torch.Tensor, torch.Tensor):
+ # ms_features are from the key encoder
+ # we only use the first one (lowest resolution), following XMem
+ if self.single_object:
+ g = masks.unsqueeze(2)
+ else:
+ g = torch.stack([masks, others], dim=2)
+
+ g = self.distributor(image, g)
+
+ batch_size, num_objects = g.shape[:2]
+ if chunk_size < 1 or chunk_size >= num_objects:
+ chunk_size = num_objects
+ fast_path = True
+ new_sensory = sensory
+ else:
+ if deep_update:
+ new_sensory = torch.empty_like(sensory)
+ else:
+ new_sensory = sensory
+ fast_path = False
+
+ # chunk-by-chunk inference
+ all_g = []
+ for i in range(0, num_objects, chunk_size):
+ if fast_path:
+ g_chunk = g
+ else:
+ g_chunk = g[:, i:i + chunk_size]
+ actual_chunk_size = g_chunk.shape[1]
+ g_chunk = g_chunk.flatten(start_dim=0, end_dim=1)
+
+ g_chunk = self.conv1(g_chunk)
+ g_chunk = self.bn1(g_chunk) # 1/2, 64
+ g_chunk = self.maxpool(g_chunk) # 1/4, 64
+ g_chunk = self.relu(g_chunk)
+
+ g_chunk = self.layer1(g_chunk) # 1/4
+ g_chunk = self.layer2(g_chunk) # 1/8
+ g_chunk = self.layer3(g_chunk) # 1/16
+
+ g_chunk = g_chunk.view(batch_size, actual_chunk_size, *g_chunk.shape[1:])
+ g_chunk = self.fuser(pix_feat, g_chunk)
+ all_g.append(g_chunk)
+ if deep_update:
+ if fast_path:
+ new_sensory = self.sensory_update(g_chunk, sensory)
+ else:
+ new_sensory[:, i:i + chunk_size] = self.sensory_update(
+ g_chunk, sensory[:, i:i + chunk_size])
+ g = torch.cat(all_g, dim=1)
+
+ return g, new_sensory
+
+ # override the default train() to freeze BN statistics
+ def train(self, mode=True):
+ self.training = False
+ for module in self.children():
+ module.train(False)
+ return self
+
+
+class PixelFeatureFuser(nn.Module):
+ def __init__(self, model_cfg: DictConfig, single_object=False):
+ super().__init__()
+ value_dim = model_cfg.value_dim
+ sensory_dim = model_cfg.sensory_dim
+ pixel_dim = model_cfg.pixel_dim
+ embed_dim = model_cfg.embed_dim
+ self.single_object = single_object
+
+ self.fuser = GroupFeatureFusionBlock(pixel_dim, value_dim, embed_dim)
+ if self.single_object:
+ self.sensory_compress = GConv2d(sensory_dim + 1, value_dim, kernel_size=1)
+ else:
+ self.sensory_compress = GConv2d(sensory_dim + 2, value_dim, kernel_size=1)
+
+ def forward(self,
+ pix_feat: torch.Tensor,
+ pixel_memory: torch.Tensor,
+ sensory_memory: torch.Tensor,
+ last_mask: torch.Tensor,
+ last_others: torch.Tensor,
+ *,
+ chunk_size: int = -1) -> torch.Tensor:
+ batch_size, num_objects = pixel_memory.shape[:2]
+
+ if self.single_object:
+ last_mask = last_mask.unsqueeze(2)
+ else:
+ last_mask = torch.stack([last_mask, last_others], dim=2)
+
+ if chunk_size < 1:
+ chunk_size = num_objects
+
+ # chunk-by-chunk inference
+ all_p16 = []
+ for i in range(0, num_objects, chunk_size):
+ sensory_readout = self.sensory_compress(
+ torch.cat([sensory_memory[:, i:i + chunk_size], last_mask[:, i:i + chunk_size]], 2))
+ p16 = pixel_memory[:, i:i + chunk_size] + sensory_readout
+ p16 = self.fuser(pix_feat, p16)
+ all_p16.append(p16)
+ p16 = torch.cat(all_p16, dim=1)
+
+ return p16
+
+
+class MaskDecoder(nn.Module):
+ def __init__(self, model_cfg: DictConfig):
+ super().__init__()
+ embed_dim = model_cfg.embed_dim
+ sensory_dim = model_cfg.sensory_dim
+ ms_image_dims = model_cfg.pixel_encoder.ms_dims
+ up_dims = model_cfg.mask_decoder.up_dims
+
+ assert embed_dim == up_dims[0]
+
+ self.sensory_update = SensoryUpdater_fullscale([up_dims[0], up_dims[1], up_dims[2], up_dims[3], up_dims[4] + 1], sensory_dim,
+ sensory_dim)
+
+ self.decoder_feat_proc = DecoderFeatureProcessor(ms_image_dims[1:], up_dims[:-1])
+ self.up_16_8 = MaskUpsampleBlock(up_dims[0], up_dims[1])
+ self.up_8_4 = MaskUpsampleBlock(up_dims[1], up_dims[2])
+ # newly add for alpha matte
+ self.up_4_2 = MaskUpsampleBlock(up_dims[2], up_dims[3])
+ self.up_2_1 = MaskUpsampleBlock(up_dims[3], up_dims[4])
+
+ self.pred_seg = nn.Conv2d(up_dims[-1], 1, kernel_size=3, padding=1)
+ self.pred_mat = nn.Conv2d(up_dims[-1], 1, kernel_size=3, padding=1)
+
+ def forward(self,
+ ms_image_feat: Iterable[torch.Tensor],
+ memory_readout: torch.Tensor,
+ sensory: torch.Tensor,
+ *,
+ chunk_size: int = -1,
+ update_sensory: bool = True,
+ seg_pass: bool = False,
+ last_mask=None,
+ sigmoid_residual=False) -> (torch.Tensor, torch.Tensor):
+
+ batch_size, num_objects = memory_readout.shape[:2]
+ f8, f4, f2, f1 = self.decoder_feat_proc(ms_image_feat[1:])
+ if chunk_size < 1 or chunk_size >= num_objects:
+ chunk_size = num_objects
+ fast_path = True
+ new_sensory = sensory
+ else:
+ if update_sensory:
+ new_sensory = torch.empty_like(sensory)
+ else:
+ new_sensory = sensory
+ fast_path = False
+
+ # chunk-by-chunk inference
+ all_logits = []
+ for i in range(0, num_objects, chunk_size):
+ if fast_path:
+ p16 = memory_readout
+ else:
+ p16 = memory_readout[:, i:i + chunk_size]
+ actual_chunk_size = p16.shape[1]
+
+ p8 = self.up_16_8(p16, f8)
+ p4 = self.up_8_4(p8, f4)
+ p2 = self.up_4_2(p4, f2)
+ p1 = self.up_2_1(p2, f1)
+ with torch.amp.autocast("cuda",enabled=False):
+ if seg_pass:
+ if last_mask is not None:
+ res = self.pred_seg(F.relu(p1.flatten(start_dim=0, end_dim=1).float()))
+ if sigmoid_residual:
+ res = (torch.sigmoid(res) - 0.5) * 2 # regularization: (-1, 1) change on last mask
+ logits = last_mask + res
+ else:
+ logits = self.pred_seg(F.relu(p1.flatten(start_dim=0, end_dim=1).float()))
+ else:
+ if last_mask is not None:
+ res = self.pred_mat(F.relu(p1.flatten(start_dim=0, end_dim=1).float()))
+ if sigmoid_residual:
+ res = (torch.sigmoid(res) - 0.5) * 2 # regularization: (-1, 1) change on last mask
+ logits = last_mask + res
+ else:
+ logits = self.pred_mat(F.relu(p1.flatten(start_dim=0, end_dim=1).float()))
+ ## SensoryUpdater_fullscale
+ if update_sensory:
+ p1 = torch.cat(
+ [p1, logits.view(batch_size, actual_chunk_size, 1, *logits.shape[-2:])], 2)
+ if fast_path:
+ new_sensory = self.sensory_update([p16, p8, p4, p2, p1], sensory)
+ else:
+ new_sensory[:,
+ i:i + chunk_size] = self.sensory_update([p16, p8, p4, p2, p1],
+ sensory[:,
+ i:i + chunk_size])
+ all_logits.append(logits)
+ logits = torch.cat(all_logits, dim=0)
+ logits = logits.view(batch_size, num_objects, *logits.shape[-2:])
+
+ return new_sensory, logits
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/model/channel_attn.py b/auto_media_publisher/modifiers/video_matting/matanyone/model/channel_attn.py
new file mode 100644
index 0000000..a2096c1
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/model/channel_attn.py
@@ -0,0 +1,39 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class CAResBlock(nn.Module):
+ def __init__(self, in_dim: int, out_dim: int, residual: bool = True):
+ super().__init__()
+ self.residual = residual
+ self.conv1 = nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1)
+ self.conv2 = nn.Conv2d(out_dim, out_dim, kernel_size=3, padding=1)
+
+ t = int((abs(math.log2(out_dim)) + 1) // 2)
+ k = t if t % 2 else t + 1
+ self.pool = nn.AdaptiveAvgPool2d(1)
+ self.conv = nn.Conv1d(1, 1, kernel_size=k, padding=(k - 1) // 2, bias=False)
+
+ if self.residual:
+ if in_dim == out_dim:
+ self.downsample = nn.Identity()
+ else:
+ self.downsample = nn.Conv2d(in_dim, out_dim, kernel_size=1)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ r = x
+ x = self.conv1(F.relu(x))
+ x = self.conv2(F.relu(x))
+
+ b, c = x.shape[:2]
+ w = self.pool(x).view(b, 1, c)
+ w = self.conv(w).transpose(-1, -2).unsqueeze(-1).sigmoid() # B*C*1*1
+
+ if self.residual:
+ x = x * w + self.downsample(r)
+ else:
+ x = x * w
+
+ return x
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/model/group_modules.py b/auto_media_publisher/modifiers/video_matting/matanyone/model/group_modules.py
new file mode 100644
index 0000000..a421ae7
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/model/group_modules.py
@@ -0,0 +1,126 @@
+from typing import Optional
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from auto_media_publisher.modifiers.video_matting.matanyone.model.channel_attn import CAResBlock
+
+def interpolate_groups(g: torch.Tensor, ratio: float, mode: str,
+ align_corners: bool) -> torch.Tensor:
+ batch_size, num_objects = g.shape[:2]
+ g = F.interpolate(g.flatten(start_dim=0, end_dim=1),
+ scale_factor=ratio,
+ mode=mode,
+ align_corners=align_corners)
+ g = g.view(batch_size, num_objects, *g.shape[1:])
+ return g
+
+
+def upsample_groups(g: torch.Tensor,
+ ratio: float = 2,
+ mode: str = 'bilinear',
+ align_corners: bool = False) -> torch.Tensor:
+ return interpolate_groups(g, ratio, mode, align_corners)
+
+
+def downsample_groups(g: torch.Tensor,
+ ratio: float = 1 / 2,
+ mode: str = 'area',
+ align_corners: bool = None) -> torch.Tensor:
+ return interpolate_groups(g, ratio, mode, align_corners)
+
+
+class GConv2d(nn.Conv2d):
+ def forward(self, g: torch.Tensor) -> torch.Tensor:
+ batch_size, num_objects = g.shape[:2]
+ g = super().forward(g.flatten(start_dim=0, end_dim=1))
+ return g.view(batch_size, num_objects, *g.shape[1:])
+
+
+class GroupResBlock(nn.Module):
+ def __init__(self, in_dim: int, out_dim: int):
+ super().__init__()
+
+ if in_dim == out_dim:
+ self.downsample = nn.Identity()
+ else:
+ self.downsample = GConv2d(in_dim, out_dim, kernel_size=1)
+
+ self.conv1 = GConv2d(in_dim, out_dim, kernel_size=3, padding=1)
+ self.conv2 = GConv2d(out_dim, out_dim, kernel_size=3, padding=1)
+
+ def forward(self, g: torch.Tensor) -> torch.Tensor:
+ out_g = self.conv1(F.relu(g))
+ out_g = self.conv2(F.relu(out_g))
+
+ g = self.downsample(g)
+
+ return out_g + g
+
+
+class MainToGroupDistributor(nn.Module):
+ def __init__(self,
+ x_transform: Optional[nn.Module] = None,
+ g_transform: Optional[nn.Module] = None,
+ method: str = 'cat',
+ reverse_order: bool = False):
+ super().__init__()
+
+ self.x_transform = x_transform
+ self.g_transform = g_transform
+ self.method = method
+ self.reverse_order = reverse_order
+
+ def forward(self, x: torch.Tensor, g: torch.Tensor, skip_expand: bool = False) -> torch.Tensor:
+ num_objects = g.shape[1]
+
+ if self.x_transform is not None:
+ x = self.x_transform(x)
+
+ if self.g_transform is not None:
+ g = self.g_transform(g)
+
+ if not skip_expand:
+ x = x.unsqueeze(1).expand(-1, num_objects, -1, -1, -1)
+ if self.method == 'cat':
+ if self.reverse_order:
+ g = torch.cat([g, x], 2)
+ else:
+ g = torch.cat([x, g], 2)
+ elif self.method == 'add':
+ g = x + g
+ elif self.method == 'mulcat':
+ g = torch.cat([x * g, g], dim=2)
+ elif self.method == 'muladd':
+ g = x * g + g
+ else:
+ raise NotImplementedError
+
+ return g
+
+
+class GroupFeatureFusionBlock(nn.Module):
+ def __init__(self, x_in_dim: int, g_in_dim: int, out_dim: int):
+ super().__init__()
+
+ x_transform = nn.Conv2d(x_in_dim, out_dim, kernel_size=1)
+ g_transform = GConv2d(g_in_dim, out_dim, kernel_size=1)
+
+ self.distributor = MainToGroupDistributor(x_transform=x_transform,
+ g_transform=g_transform,
+ method='add')
+ self.block1 = CAResBlock(out_dim, out_dim)
+ self.block2 = CAResBlock(out_dim, out_dim)
+
+ def forward(self, x: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
+ batch_size, num_objects = g.shape[:2]
+
+ g = self.distributor(x, g)
+
+ g = g.flatten(start_dim=0, end_dim=1)
+
+ g = self.block1(g)
+ g = self.block2(g)
+
+ g = g.view(batch_size, num_objects, *g.shape[1:])
+
+ return g
\ No newline at end of file
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/model/matanyone.py b/auto_media_publisher/modifiers/video_matting/matanyone/model/matanyone.py
new file mode 100644
index 0000000..11240f1
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/model/matanyone.py
@@ -0,0 +1,333 @@
+from typing import List, Dict, Iterable, Tuple
+import logging
+from omegaconf import DictConfig
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from omegaconf import OmegaConf
+from huggingface_hub import PyTorchModelHubMixin
+
+from auto_media_publisher.modifiers.video_matting.matanyone.model.big_modules import PixelEncoder, UncertPred, KeyProjection, MaskEncoder, PixelFeatureFuser, MaskDecoder
+from auto_media_publisher.modifiers.video_matting.matanyone.model.aux_modules import AuxComputer
+from auto_media_publisher.modifiers.video_matting.matanyone.model.utils.memory_utils import get_affinity, readout
+from auto_media_publisher.modifiers.video_matting.matanyone.model.transformer.object_transformer import QueryTransformer
+from auto_media_publisher.modifiers.video_matting.matanyone.model.transformer.object_summarizer import ObjectSummarizer
+from auto_media_publisher.modifiers.video_matting.matanyone.utils.tensor_utils import aggregate
+
+log = logging.getLogger()
+class MatAnyone(nn.Module,
+ PyTorchModelHubMixin,
+ library_name="matanyone",
+ repo_url="https://github.com/pq-yang/MatAnyone",
+ coders={
+ DictConfig: (
+ lambda x: OmegaConf.to_container(x),
+ lambda data: OmegaConf.create(data),
+ )
+ },
+ ):
+
+ def __init__(self, cfg: DictConfig, *, single_object=False):
+ super().__init__()
+ self.cfg = cfg
+ model_cfg = cfg.model
+ self.ms_dims = model_cfg.pixel_encoder.ms_dims
+ self.key_dim = model_cfg.key_dim
+ self.value_dim = model_cfg.value_dim
+ self.sensory_dim = model_cfg.sensory_dim
+ self.pixel_dim = model_cfg.pixel_dim
+ self.embed_dim = model_cfg.embed_dim
+ self.single_object = single_object
+
+ log.info(f'Single object: {self.single_object}')
+
+ self.pixel_encoder = PixelEncoder(model_cfg)
+ self.pix_feat_proj = nn.Conv2d(self.ms_dims[0], self.pixel_dim, kernel_size=1)
+ self.key_proj = KeyProjection(model_cfg)
+ self.mask_encoder = MaskEncoder(model_cfg, single_object=single_object)
+ self.mask_decoder = MaskDecoder(model_cfg)
+ self.pixel_fuser = PixelFeatureFuser(model_cfg, single_object=single_object)
+ self.object_transformer = QueryTransformer(model_cfg)
+ self.object_summarizer = ObjectSummarizer(model_cfg)
+ self.aux_computer = AuxComputer(cfg)
+ self.temp_sparity = UncertPred(model_cfg)
+
+ self.register_buffer("pixel_mean", torch.Tensor(model_cfg.pixel_mean).view(-1, 1, 1), False)
+ self.register_buffer("pixel_std", torch.Tensor(model_cfg.pixel_std).view(-1, 1, 1), False)
+
+ def _get_others(self, masks: torch.Tensor) -> torch.Tensor:
+ # for each object, return the sum of masks of all other objects
+ if self.single_object:
+ return None
+
+ num_objects = masks.shape[1]
+ if num_objects >= 1:
+ others = (masks.sum(dim=1, keepdim=True) - masks).clamp(0, 1)
+ else:
+ others = torch.zeros_like(masks)
+ return others
+
+ def pred_uncertainty(self, last_pix_feat: torch.Tensor, cur_pix_feat: torch.Tensor, last_mask: torch.Tensor, mem_val_diff:torch.Tensor):
+ logits = self.temp_sparity(last_frame_feat=last_pix_feat,
+ cur_frame_feat=cur_pix_feat,
+ last_mask=last_mask,
+ mem_val_diff=mem_val_diff)
+
+ prob = torch.sigmoid(logits)
+ mask = (prob > 0) + 0
+
+ uncert_output = {"logits": logits,
+ "prob": prob,
+ "mask": mask}
+
+ return uncert_output
+
+ def encode_image(self, image: torch.Tensor, seq_length=None, last_feats=None) -> (Iterable[torch.Tensor], torch.Tensor): # type: ignore
+ image = (image - self.pixel_mean) / self.pixel_std
+ ms_image_feat = self.pixel_encoder(image, seq_length) # f16, f8, f4, f2, f1
+ return ms_image_feat, self.pix_feat_proj(ms_image_feat[0])
+
+ def encode_mask(
+ self,
+ image: torch.Tensor,
+ ms_features: List[torch.Tensor],
+ sensory: torch.Tensor,
+ masks: torch.Tensor,
+ *,
+ deep_update: bool = True,
+ chunk_size: int = -1,
+ need_weights: bool = False) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ image = (image - self.pixel_mean) / self.pixel_std
+ others = self._get_others(masks)
+ mask_value, new_sensory = self.mask_encoder(image,
+ ms_features,
+ sensory,
+ masks,
+ others,
+ deep_update=deep_update,
+ chunk_size=chunk_size)
+ object_summaries, object_logits = self.object_summarizer(masks, mask_value, need_weights)
+ return mask_value, new_sensory, object_summaries, object_logits
+
+ def transform_key(self,
+ final_pix_feat: torch.Tensor,
+ *,
+ need_sk: bool = True,
+ need_ek: bool = True) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ key, shrinkage, selection = self.key_proj(final_pix_feat, need_s=need_sk, need_e=need_ek)
+ return key, shrinkage, selection
+
+ # Used in training only.
+ # This step is replaced by MemoryManager in test time
+ def read_memory(self, query_key: torch.Tensor, query_selection: torch.Tensor,
+ memory_key: torch.Tensor, memory_shrinkage: torch.Tensor,
+ msk_value: torch.Tensor, obj_memory: torch.Tensor, pix_feat: torch.Tensor,
+ sensory: torch.Tensor, last_mask: torch.Tensor,
+ selector: torch.Tensor, uncert_output=None, seg_pass=False,
+ last_pix_feat=None, last_pred_mask=None) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+ """
+ query_key : B * CK * H * W
+ query_selection : B * CK * H * W
+ memory_key : B * CK * T * H * W
+ memory_shrinkage: B * 1 * T * H * W
+ msk_value : B * num_objects * CV * T * H * W
+ obj_memory : B * num_objects * T * num_summaries * C
+ pixel_feature : B * C * H * W
+ """
+ batch_size, num_objects = msk_value.shape[:2]
+
+ uncert_mask = uncert_output["mask"] if uncert_output is not None else None
+
+ # read using visual attention
+ with torch.amp.autocast("cuda",enabled=False):
+ affinity = get_affinity(memory_key.float(), memory_shrinkage.float(), query_key.float(),
+ query_selection.float(), uncert_mask=uncert_mask)
+
+ msk_value = msk_value.flatten(start_dim=1, end_dim=2).float()
+
+ # B * (num_objects*CV) * H * W
+ pixel_readout = readout(affinity, msk_value, uncert_mask)
+ pixel_readout = pixel_readout.view(batch_size, num_objects, self.value_dim,
+ *pixel_readout.shape[-2:])
+
+ uncert_output = self.pred_uncertainty(last_pix_feat, pix_feat, last_pred_mask, pixel_readout[:,0]-msk_value[:,:,-1])
+ uncert_prob = uncert_output["prob"].unsqueeze(1) # b n 1 h w
+ pixel_readout = pixel_readout*uncert_prob + msk_value[:,:,-1].unsqueeze(1)*(1-uncert_prob)
+
+ pixel_readout = self.pixel_fusion(pix_feat, pixel_readout, sensory, last_mask)
+
+
+ # read from query transformer
+ mem_readout, aux_features = self.readout_query(pixel_readout, obj_memory, selector=selector, seg_pass=seg_pass)
+
+ aux_output = {
+ 'sensory': sensory,
+ 'q_logits': aux_features['logits'] if aux_features else None,
+ 'attn_mask': aux_features['attn_mask'] if aux_features else None,
+ }
+
+ return mem_readout, aux_output, uncert_output
+
+ def read_first_frame_memory(self, pixel_readout,
+ obj_memory: torch.Tensor, pix_feat: torch.Tensor,
+ sensory: torch.Tensor, last_mask: torch.Tensor,
+ selector: torch.Tensor, seg_pass=False) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+ """
+ query_key : B * CK * H * W
+ query_selection : B * CK * H * W
+ memory_key : B * CK * T * H * W
+ memory_shrinkage: B * 1 * T * H * W
+ msk_value : B * num_objects * CV * T * H * W
+ obj_memory : B * num_objects * T * num_summaries * C
+ pixel_feature : B * C * H * W
+ """
+
+ pixel_readout = self.pixel_fusion(pix_feat, pixel_readout, sensory, last_mask)
+
+ # read from query transformer
+ mem_readout, aux_features = self.readout_query(pixel_readout, obj_memory, selector=selector, seg_pass=seg_pass)
+
+ aux_output = {
+ 'sensory': sensory,
+ 'q_logits': aux_features['logits'] if aux_features else None,
+ 'attn_mask': aux_features['attn_mask'] if aux_features else None,
+ }
+
+ return mem_readout, aux_output
+
+ def pixel_fusion(self,
+ pix_feat: torch.Tensor,
+ pixel: torch.Tensor,
+ sensory: torch.Tensor,
+ last_mask: torch.Tensor,
+ *,
+ chunk_size: int = -1) -> torch.Tensor:
+ last_mask = F.interpolate(last_mask, size=sensory.shape[-2:], mode='area')
+ last_others = self._get_others(last_mask)
+ fused = self.pixel_fuser(pix_feat,
+ pixel,
+ sensory,
+ last_mask,
+ last_others,
+ chunk_size=chunk_size)
+ return fused
+
+ def readout_query(self,
+ pixel_readout,
+ obj_memory,
+ *,
+ selector=None,
+ need_weights=False,
+ seg_pass=False) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+ return self.object_transformer(pixel_readout,
+ obj_memory,
+ selector=selector,
+ need_weights=need_weights,
+ seg_pass=seg_pass)
+
+ def segment(self,
+ ms_image_feat: List[torch.Tensor],
+ memory_readout: torch.Tensor,
+ sensory: torch.Tensor,
+ *,
+ selector: bool = None,
+ chunk_size: int = -1,
+ update_sensory: bool = True,
+ seg_pass: bool = False,
+ clamp_mat: bool = True,
+ last_mask=None,
+ sigmoid_residual=False,
+ seg_mat=False) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ multi_scale_features is from the key encoder for skip-connection
+ memory_readout is from working/long-term memory
+ sensory is the sensory memory
+ last_mask is the mask from the last frame, supplementing sensory memory
+ selector is 1 if an object exists, and 0 otherwise. We use it to filter padded objects
+ during training.
+ """
+ #### use mat head for seg data
+ if seg_mat:
+ assert seg_pass
+ seg_pass = False
+ ####
+ sensory, logits = self.mask_decoder(ms_image_feat,
+ memory_readout,
+ sensory,
+ chunk_size=chunk_size,
+ update_sensory=update_sensory,
+ seg_pass = seg_pass,
+ last_mask=last_mask,
+ sigmoid_residual=sigmoid_residual)
+ if seg_pass:
+ prob = torch.sigmoid(logits)
+ if selector is not None:
+ prob = prob * selector
+
+ # Softmax over all objects[]
+ logits = aggregate(prob, dim=1)
+ prob = F.softmax(logits, dim=1)
+ else:
+ if clamp_mat:
+ logits = logits.clamp(0.0, 1.0)
+ logits = torch.cat([torch.prod(1 - logits, dim=1, keepdim=True), logits], 1)
+ prob = logits
+
+ return sensory, logits, prob
+
+ def compute_aux(self, pix_feat: torch.Tensor, aux_inputs: Dict[str, torch.Tensor],
+ selector: torch.Tensor, seg_pass=False) -> Dict[str, torch.Tensor]:
+ return self.aux_computer(pix_feat, aux_inputs, selector, seg_pass=seg_pass)
+
+ def forward(self, *args, **kwargs):
+ raise NotImplementedError
+
+ def load_weights(self, src_dict, init_as_zero_if_needed=False) -> None:
+ if not self.single_object:
+ # Map single-object weight to multi-object weight (4->5 out channels in conv1)
+ for k in list(src_dict.keys()):
+ if k == 'mask_encoder.conv1.weight':
+ if src_dict[k].shape[1] == 4:
+ log.info(f'Converting {k} from single object to multiple objects.')
+ pads = torch.zeros((64, 1, 7, 7), device=src_dict[k].device)
+ if not init_as_zero_if_needed:
+ nn.init.orthogonal_(pads)
+ log.info(f'Randomly initialized padding for {k}.')
+ else:
+ log.info(f'Zero-initialized padding for {k}.')
+ src_dict[k] = torch.cat([src_dict[k], pads], 1)
+ elif k == 'pixel_fuser.sensory_compress.weight':
+ if src_dict[k].shape[1] == self.sensory_dim + 1:
+ log.info(f'Converting {k} from single object to multiple objects.')
+ pads = torch.zeros((self.value_dim, 1, 1, 1), device=src_dict[k].device)
+ if not init_as_zero_if_needed:
+ nn.init.orthogonal_(pads)
+ log.info(f'Randomly initialized padding for {k}.')
+ else:
+ log.info(f'Zero-initialized padding for {k}.')
+ src_dict[k] = torch.cat([src_dict[k], pads], 1)
+ elif self.single_object:
+ """
+ If the model is multiple-object and we are training in single-object,
+ we strip the last channel of conv1.
+ This is not supposed to happen in standard training except when users are trying to
+ finetune a trained model with single object datasets.
+ """
+ if src_dict['mask_encoder.conv1.weight'].shape[1] == 5:
+ log.warning('Converting mask_encoder.conv1.weight from multiple objects to single object.'
+ 'This is not supposed to happen in standard training.')
+ src_dict['mask_encoder.conv1.weight'] = src_dict['mask_encoder.conv1.weight'][:, :-1]
+ src_dict['pixel_fuser.sensory_compress.weight'] = src_dict['pixel_fuser.sensory_compress.weight'][:, :-1]
+
+ for k in src_dict:
+ if k not in self.state_dict():
+ log.info(f'Key {k} found in src_dict but not in self.state_dict()!!!')
+ for k in self.state_dict():
+ if k not in src_dict:
+ log.info(f'Key {k} found in self.state_dict() but not in src_dict!!!')
+
+ self.load_state_dict(src_dict, strict=False)
+
+ @property
+ def device(self) -> torch.device:
+ return self.pixel_mean.device
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/model/modules.py b/auto_media_publisher/modifiers/video_matting/matanyone/model/modules.py
new file mode 100644
index 0000000..9c0df7b
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/model/modules.py
@@ -0,0 +1,149 @@
+from typing import List, Iterable
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from auto_media_publisher.modifiers.video_matting.matanyone.model.group_modules import MainToGroupDistributor, GroupResBlock, upsample_groups, GConv2d, downsample_groups
+
+
+class UpsampleBlock(nn.Module):
+ def __init__(self, in_dim: int, out_dim: int, scale_factor: int = 2):
+ super().__init__()
+ self.out_conv = ResBlock(in_dim, out_dim)
+ self.scale_factor = scale_factor
+
+ def forward(self, in_g: torch.Tensor, skip_f: torch.Tensor) -> torch.Tensor:
+ g = F.interpolate(in_g,
+ scale_factor=self.scale_factor,
+ mode='bilinear')
+ g = self.out_conv(g)
+ g = g + skip_f
+ return g
+
+class MaskUpsampleBlock(nn.Module):
+ def __init__(self, in_dim: int, out_dim: int, scale_factor: int = 2):
+ super().__init__()
+ self.distributor = MainToGroupDistributor(method='add')
+ self.out_conv = GroupResBlock(in_dim, out_dim)
+ self.scale_factor = scale_factor
+
+ def forward(self, in_g: torch.Tensor, skip_f: torch.Tensor) -> torch.Tensor:
+ g = upsample_groups(in_g, ratio=self.scale_factor)
+ g = self.distributor(skip_f, g)
+ g = self.out_conv(g)
+ return g
+
+
+class DecoderFeatureProcessor(nn.Module):
+ def __init__(self, decoder_dims: List[int], out_dims: List[int]):
+ super().__init__()
+ self.transforms = nn.ModuleList([
+ nn.Conv2d(d_dim, p_dim, kernel_size=1) for d_dim, p_dim in zip(decoder_dims, out_dims)
+ ])
+
+ def forward(self, multi_scale_features: Iterable[torch.Tensor]) -> List[torch.Tensor]:
+ outputs = [func(x) for x, func in zip(multi_scale_features, self.transforms)]
+ return outputs
+
+
+# @torch.jit.script
+def _recurrent_update(h: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
+ # h: batch_size * num_objects * hidden_dim * h * w
+ # values: batch_size * num_objects * (hidden_dim*3) * h * w
+ dim = values.shape[2] // 3
+ forget_gate = torch.sigmoid(values[:, :, :dim])
+ update_gate = torch.sigmoid(values[:, :, dim:dim * 2])
+ new_value = torch.tanh(values[:, :, dim * 2:])
+ new_h = forget_gate * h * (1 - update_gate) + update_gate * new_value
+ return new_h
+
+
+class SensoryUpdater_fullscale(nn.Module):
+ # Used in the decoder, multi-scale feature + GRU
+ def __init__(self, g_dims: List[int], mid_dim: int, sensory_dim: int):
+ super().__init__()
+ self.g16_conv = GConv2d(g_dims[0], mid_dim, kernel_size=1)
+ self.g8_conv = GConv2d(g_dims[1], mid_dim, kernel_size=1)
+ self.g4_conv = GConv2d(g_dims[2], mid_dim, kernel_size=1)
+ self.g2_conv = GConv2d(g_dims[3], mid_dim, kernel_size=1)
+ self.g1_conv = GConv2d(g_dims[4], mid_dim, kernel_size=1)
+
+ self.transform = GConv2d(mid_dim + sensory_dim, sensory_dim * 3, kernel_size=3, padding=1)
+
+ nn.init.xavier_normal_(self.transform.weight)
+
+ def forward(self, g: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
+ g = self.g16_conv(g[0]) + self.g8_conv(downsample_groups(g[1], ratio=1/2)) + \
+ self.g4_conv(downsample_groups(g[2], ratio=1/4)) + \
+ self.g2_conv(downsample_groups(g[3], ratio=1/8)) + \
+ self.g1_conv(downsample_groups(g[4], ratio=1/16))
+
+ with torch.amp.autocast("cuda",enabled=False):
+ g = g.float()
+ h = h.float()
+ values = self.transform(torch.cat([g, h], dim=2))
+ new_h = _recurrent_update(h, values)
+
+ return new_h
+
+class SensoryUpdater(nn.Module):
+ # Used in the decoder, multi-scale feature + GRU
+ def __init__(self, g_dims: List[int], mid_dim: int, sensory_dim: int):
+ super().__init__()
+ self.g16_conv = GConv2d(g_dims[0], mid_dim, kernel_size=1)
+ self.g8_conv = GConv2d(g_dims[1], mid_dim, kernel_size=1)
+ self.g4_conv = GConv2d(g_dims[2], mid_dim, kernel_size=1)
+
+ self.transform = GConv2d(mid_dim + sensory_dim, sensory_dim * 3, kernel_size=3, padding=1)
+
+ nn.init.xavier_normal_(self.transform.weight)
+
+ def forward(self, g: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
+ g = self.g16_conv(g[0]) + self.g8_conv(downsample_groups(g[1], ratio=1/2)) + \
+ self.g4_conv(downsample_groups(g[2], ratio=1/4))
+
+ with torch.amp.autocast("cuda",enabled=False):
+ g = g.float()
+ h = h.float()
+ values = self.transform(torch.cat([g, h], dim=2))
+ new_h = _recurrent_update(h, values)
+
+ return new_h
+
+
+class SensoryDeepUpdater(nn.Module):
+ def __init__(self, f_dim: int, sensory_dim: int):
+ super().__init__()
+ self.transform = GConv2d(f_dim + sensory_dim, sensory_dim * 3, kernel_size=3, padding=1)
+
+ nn.init.xavier_normal_(self.transform.weight)
+
+ def forward(self, g: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
+ with torch.amp.autocast("cuda",enabled=False):
+ g = g.float()
+ h = h.float()
+ values = self.transform(torch.cat([g, h], dim=2))
+ new_h = _recurrent_update(h, values)
+
+ return new_h
+
+
+class ResBlock(nn.Module):
+ def __init__(self, in_dim: int, out_dim: int):
+ super().__init__()
+
+ if in_dim == out_dim:
+ self.downsample = nn.Identity()
+ else:
+ self.downsample = nn.Conv2d(in_dim, out_dim, kernel_size=1)
+
+ self.conv1 = nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1)
+ self.conv2 = nn.Conv2d(out_dim, out_dim, kernel_size=3, padding=1)
+
+ def forward(self, g: torch.Tensor) -> torch.Tensor:
+ out_g = self.conv1(F.relu(g))
+ out_g = self.conv2(F.relu(out_g))
+
+ g = self.downsample(g)
+
+ return out_g + g
\ No newline at end of file
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/__init__.py b/auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/__init__.py
similarity index 100%
copy from auto_media_publisher/modifiers/video_matting/matting_model/__init__.py
copy to auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/__init__.py
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/object_summarizer.py b/auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/object_summarizer.py
new file mode 100644
index 0000000..4e3c269
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/object_summarizer.py
@@ -0,0 +1,89 @@
+from typing import Optional
+from omegaconf import DictConfig
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from auto_media_publisher.modifiers.video_matting.matanyone.model.transformer.positional_encoding import PositionalEncoding
+
+
+# @torch.jit.script
+def _weighted_pooling(masks: torch.Tensor, value: torch.Tensor,
+ logits: torch.Tensor) -> (torch.Tensor, torch.Tensor):
+ # value: B*num_objects*H*W*value_dim
+ # logits: B*num_objects*H*W*num_summaries
+ # masks: B*num_objects*H*W*num_summaries: 1 if allowed
+ weights = logits.sigmoid() * masks
+ # B*num_objects*num_summaries*value_dim
+ sums = torch.einsum('bkhwq,bkhwc->bkqc', weights, value)
+ # B*num_objects*H*W*num_summaries -> B*num_objects*num_summaries*1
+ area = weights.flatten(start_dim=2, end_dim=3).sum(2).unsqueeze(-1)
+
+ # B*num_objects*num_summaries*value_dim
+ return sums, area
+
+
+class ObjectSummarizer(nn.Module):
+ def __init__(self, model_cfg: DictConfig):
+ super().__init__()
+
+ this_cfg = model_cfg.object_summarizer
+ self.value_dim = model_cfg.value_dim
+ self.embed_dim = this_cfg.embed_dim
+ self.num_summaries = this_cfg.num_summaries
+ self.add_pe = this_cfg.add_pe
+ self.pixel_pe_scale = model_cfg.pixel_pe_scale
+ self.pixel_pe_temperature = model_cfg.pixel_pe_temperature
+
+ if self.add_pe:
+ self.pos_enc = PositionalEncoding(self.embed_dim,
+ scale=self.pixel_pe_scale,
+ temperature=self.pixel_pe_temperature)
+
+ self.input_proj = nn.Linear(self.value_dim, self.embed_dim)
+ self.feature_pred = nn.Sequential(
+ nn.Linear(self.embed_dim, self.embed_dim),
+ nn.ReLU(inplace=True),
+ nn.Linear(self.embed_dim, self.embed_dim),
+ )
+ self.weights_pred = nn.Sequential(
+ nn.Linear(self.embed_dim, self.embed_dim),
+ nn.ReLU(inplace=True),
+ nn.Linear(self.embed_dim, self.num_summaries),
+ )
+
+ def forward(self,
+ masks: torch.Tensor,
+ value: torch.Tensor,
+ need_weights: bool = False) -> (torch.Tensor, Optional[torch.Tensor]):
+ # masks: B*num_objects*(H0)*(W0)
+ # value: B*num_objects*value_dim*H*W
+ # -> B*num_objects*H*W*value_dim
+ h, w = value.shape[-2:]
+ masks = F.interpolate(masks, size=(h, w), mode='area')
+ masks = masks.unsqueeze(-1)
+ inv_masks = 1 - masks
+ repeated_masks = torch.cat([
+ masks.expand(-1, -1, -1, -1, self.num_summaries // 2),
+ inv_masks.expand(-1, -1, -1, -1, self.num_summaries // 2),
+ ],
+ dim=-1)
+
+ value = value.permute(0, 1, 3, 4, 2)
+ value = self.input_proj(value)
+ if self.add_pe:
+ pe = self.pos_enc(value)
+ value = value + pe
+
+ with torch.amp.autocast("cuda",enabled=False):
+ value = value.float()
+ feature = self.feature_pred(value)
+ logits = self.weights_pred(value)
+ sums, area = _weighted_pooling(repeated_masks, feature, logits)
+
+ summaries = torch.cat([sums, area], dim=-1)
+
+ if need_weights:
+ return summaries, logits
+ else:
+ return summaries, None
\ No newline at end of file
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/object_transformer.py b/auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/object_transformer.py
new file mode 100644
index 0000000..a38bfbb
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/object_transformer.py
@@ -0,0 +1,206 @@
+from typing import Dict, Optional
+from omegaconf import DictConfig
+
+import torch
+import torch.nn as nn
+from auto_media_publisher.modifiers.video_matting.matanyone.model.group_modules import GConv2d
+from auto_media_publisher.modifiers.video_matting.matanyone.utils.tensor_utils import aggregate
+from auto_media_publisher.modifiers.video_matting.matanyone.model.transformer.positional_encoding import PositionalEncoding
+from auto_media_publisher.modifiers.video_matting.matanyone.model.transformer.transformer_layers import CrossAttention, SelfAttention, FFN, PixelFFN
+
+
+class QueryTransformerBlock(nn.Module):
+ def __init__(self, model_cfg: DictConfig):
+ super().__init__()
+
+ this_cfg = model_cfg.object_transformer
+ self.embed_dim = this_cfg.embed_dim
+ self.num_heads = this_cfg.num_heads
+ self.num_queries = this_cfg.num_queries
+ self.ff_dim = this_cfg.ff_dim
+
+ self.read_from_pixel = CrossAttention(self.embed_dim,
+ self.num_heads,
+ add_pe_to_qkv=this_cfg.read_from_pixel.add_pe_to_qkv)
+ self.self_attn = SelfAttention(self.embed_dim,
+ self.num_heads,
+ add_pe_to_qkv=this_cfg.query_self_attention.add_pe_to_qkv)
+ self.ffn = FFN(self.embed_dim, self.ff_dim)
+ self.read_from_query = CrossAttention(self.embed_dim,
+ self.num_heads,
+ add_pe_to_qkv=this_cfg.read_from_query.add_pe_to_qkv,
+ norm=this_cfg.read_from_query.output_norm)
+ self.pixel_ffn = PixelFFN(self.embed_dim)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ pixel: torch.Tensor,
+ query_pe: torch.Tensor,
+ pixel_pe: torch.Tensor,
+ attn_mask: torch.Tensor,
+ need_weights: bool = False) -> (torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor):
+ # x: (bs*num_objects)*num_queries*embed_dim
+ # pixel: bs*num_objects*C*H*W
+ # query_pe: (bs*num_objects)*num_queries*embed_dim
+ # pixel_pe: (bs*num_objects)*(H*W)*C
+ # attn_mask: (bs*num_objects*num_heads)*num_queries*(H*W)
+
+ # bs*num_objects*C*H*W -> (bs*num_objects)*(H*W)*C
+ pixel_flat = pixel.flatten(3, 4).flatten(0, 1).transpose(1, 2).contiguous()
+ x, q_weights = self.read_from_pixel(x,
+ pixel_flat,
+ query_pe,
+ pixel_pe,
+ attn_mask=attn_mask,
+ need_weights=need_weights)
+ x = self.self_attn(x, query_pe)
+ x = self.ffn(x)
+
+ pixel_flat, p_weights = self.read_from_query(pixel_flat,
+ x,
+ pixel_pe,
+ query_pe,
+ need_weights=need_weights)
+ pixel = self.pixel_ffn(pixel, pixel_flat)
+
+ if need_weights:
+ bs, num_objects, _, h, w = pixel.shape
+ q_weights = q_weights.view(bs, num_objects, self.num_heads, self.num_queries, h, w)
+ p_weights = p_weights.transpose(2, 3).view(bs, num_objects, self.num_heads,
+ self.num_queries, h, w)
+
+ return x, pixel, q_weights, p_weights
+
+
+class QueryTransformer(nn.Module):
+ def __init__(self, model_cfg: DictConfig):
+ super().__init__()
+
+ this_cfg = model_cfg.object_transformer
+ self.value_dim = model_cfg.value_dim
+ self.embed_dim = this_cfg.embed_dim
+ self.num_heads = this_cfg.num_heads
+ self.num_queries = this_cfg.num_queries
+
+ # query initialization and embedding
+ self.query_init = nn.Embedding(self.num_queries, self.embed_dim)
+ self.query_emb = nn.Embedding(self.num_queries, self.embed_dim)
+
+ # projection from object summaries to query initialization and embedding
+ self.summary_to_query_init = nn.Linear(self.embed_dim, self.embed_dim)
+ self.summary_to_query_emb = nn.Linear(self.embed_dim, self.embed_dim)
+
+ self.pixel_pe_scale = model_cfg.pixel_pe_scale
+ self.pixel_pe_temperature = model_cfg.pixel_pe_temperature
+ self.pixel_init_proj = GConv2d(self.embed_dim, self.embed_dim, kernel_size=1)
+ self.pixel_emb_proj = GConv2d(self.embed_dim, self.embed_dim, kernel_size=1)
+ self.spatial_pe = PositionalEncoding(self.embed_dim,
+ scale=self.pixel_pe_scale,
+ temperature=self.pixel_pe_temperature,
+ channel_last=False,
+ transpose_output=True)
+
+ # transformer blocks
+ self.num_blocks = this_cfg.num_blocks
+ self.blocks = nn.ModuleList(
+ QueryTransformerBlock(model_cfg) for _ in range(self.num_blocks))
+ self.mask_pred = nn.ModuleList(
+ nn.Sequential(nn.ReLU(), GConv2d(self.embed_dim, 1, kernel_size=1))
+ for _ in range(self.num_blocks + 1))
+
+ self.act = nn.ReLU(inplace=True)
+
+ def forward(self,
+ pixel: torch.Tensor,
+ obj_summaries: torch.Tensor,
+ selector: Optional[torch.Tensor] = None,
+ need_weights: bool = False,
+ seg_pass=False) -> (torch.Tensor, Dict[str, torch.Tensor]):
+ # pixel: B*num_objects*embed_dim*H*W
+ # obj_summaries: B*num_objects*T*num_queries*embed_dim
+ T = obj_summaries.shape[2]
+ bs, num_objects, _, H, W = pixel.shape
+
+ # normalize object values
+ # the last channel is the cumulative area of the object
+ obj_summaries = obj_summaries.view(bs * num_objects, T, self.num_queries,
+ self.embed_dim + 1)
+ # sum over time
+ # during inference, T=1 as we already did streaming average in memory_manager
+ obj_sums = obj_summaries[:, :, :, :-1].sum(dim=1)
+ obj_area = obj_summaries[:, :, :, -1:].sum(dim=1)
+ obj_values = obj_sums / (obj_area + 1e-4)
+ obj_init = self.summary_to_query_init(obj_values)
+ obj_emb = self.summary_to_query_emb(obj_values)
+
+ # positional embeddings for object queries
+ query = self.query_init.weight.unsqueeze(0).expand(bs * num_objects, -1, -1) + obj_init
+ query_emb = self.query_emb.weight.unsqueeze(0).expand(bs * num_objects, -1, -1) + obj_emb
+
+ # positional embeddings for pixel features
+ pixel_init = self.pixel_init_proj(pixel)
+ pixel_emb = self.pixel_emb_proj(pixel)
+ pixel_pe = self.spatial_pe(pixel.flatten(0, 1))
+ pixel_emb = pixel_emb.flatten(3, 4).flatten(0, 1).transpose(1, 2).contiguous()
+ pixel_pe = pixel_pe.flatten(1, 2) + pixel_emb
+
+ pixel = pixel_init
+
+ # run the transformer
+ aux_features = {'logits': []}
+
+ # first aux output
+ aux_logits = self.mask_pred[0](pixel).squeeze(2)
+ attn_mask = self._get_aux_mask(aux_logits, selector, seg_pass=seg_pass)
+ aux_features['logits'].append(aux_logits)
+ for i in range(self.num_blocks):
+ query, pixel, q_weights, p_weights = self.blocks[i](query,
+ pixel,
+ query_emb,
+ pixel_pe,
+ attn_mask,
+ need_weights=need_weights)
+
+ if self.training or i <= self.num_blocks - 1 or need_weights:
+ aux_logits = self.mask_pred[i + 1](pixel).squeeze(2)
+ attn_mask = self._get_aux_mask(aux_logits, selector, seg_pass=seg_pass)
+ aux_features['logits'].append(aux_logits)
+
+ aux_features['q_weights'] = q_weights # last layer only
+ aux_features['p_weights'] = p_weights # last layer only
+
+ if self.training:
+ # no need to save all heads
+ aux_features['attn_mask'] = attn_mask.view(bs, num_objects, self.num_heads,
+ self.num_queries, H, W)[:, :, 0]
+
+ return pixel, aux_features
+
+ def _get_aux_mask(self, logits: torch.Tensor, selector: torch.Tensor, seg_pass=False) -> torch.Tensor:
+ # logits: batch_size*num_objects*H*W
+ # selector: batch_size*num_objects*1*1
+ # returns a mask of shape (batch_size*num_objects*num_heads)*num_queries*(H*W)
+ # where True means the attention is blocked
+
+ if selector is None:
+ prob = logits.sigmoid()
+ else:
+ prob = logits.sigmoid() * selector
+ logits = aggregate(prob, dim=1)
+
+ is_foreground = (logits[:, 1:] >= logits.max(dim=1, keepdim=True)[0])
+ foreground_mask = is_foreground.bool().flatten(start_dim=2)
+ inv_foreground_mask = ~foreground_mask
+ inv_background_mask = foreground_mask
+
+ aux_foreground_mask = inv_foreground_mask.unsqueeze(2).unsqueeze(2).repeat(
+ 1, 1, self.num_heads, self.num_queries // 2, 1).flatten(start_dim=0, end_dim=2)
+ aux_background_mask = inv_background_mask.unsqueeze(2).unsqueeze(2).repeat(
+ 1, 1, self.num_heads, self.num_queries // 2, 1).flatten(start_dim=0, end_dim=2)
+
+ aux_mask = torch.cat([aux_foreground_mask, aux_background_mask], dim=1)
+
+ aux_mask[torch.where(aux_mask.sum(-1) == aux_mask.shape[-1])] = False
+
+ return aux_mask
\ No newline at end of file
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/positional_encoding.py b/auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/positional_encoding.py
new file mode 100644
index 0000000..6c15bb7
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/positional_encoding.py
@@ -0,0 +1,108 @@
+# Reference:
+# https://github.com/facebookresearch/Mask2Former/blob/main/mask2former/modeling/transformer_decoder/position_encoding.py
+# https://github.com/tatp22/multidim-positional-encoding/blob/master/positional_encodings/torch_encodings.py
+
+import math
+
+import numpy as np
+import torch
+from torch import nn
+
+
+def get_emb(sin_inp: torch.Tensor) -> torch.Tensor:
+ """
+ Gets a base embedding for one dimension with sin and cos intertwined
+ """
+ emb = torch.stack((sin_inp.sin(), sin_inp.cos()), dim=-1)
+ return torch.flatten(emb, -2, -1)
+
+
+class PositionalEncoding(nn.Module):
+ def __init__(self,
+ dim: int,
+ scale: float = math.pi * 2,
+ temperature: float = 10000,
+ normalize: bool = True,
+ channel_last: bool = True,
+ transpose_output: bool = False):
+ super().__init__()
+ dim = int(np.ceil(dim / 4) * 2)
+ self.dim = dim
+ inv_freq = 1.0 / (temperature**(torch.arange(0, dim, 2).float() / dim))
+ self.register_buffer("inv_freq", inv_freq)
+ self.normalize = normalize
+ self.scale = scale
+ self.eps = 1e-6
+ self.channel_last = channel_last
+ self.transpose_output = transpose_output
+
+ self.cached_penc = None # the cache is irrespective of the number of objects
+
+ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+ """
+ :param tensor: A 4/5d tensor of size
+ channel_last=True: (batch_size, h, w, c) or (batch_size, k, h, w, c)
+ channel_last=False: (batch_size, c, h, w) or (batch_size, k, c, h, w)
+ :return: positional encoding tensor that has the same shape as the input if the input is 4d
+ if the input is 5d, the output is broadcastable along the k-dimension
+ """
+ if len(tensor.shape) != 4 and len(tensor.shape) != 5:
+ raise RuntimeError(f'The input tensor has to be 4/5d, got {tensor.shape}!')
+
+ if len(tensor.shape) == 5:
+ # take a sample from the k dimension
+ num_objects = tensor.shape[1]
+ tensor = tensor[:, 0]
+ else:
+ num_objects = None
+
+ if self.channel_last:
+ batch_size, h, w, c = tensor.shape
+ else:
+ batch_size, c, h, w = tensor.shape
+
+ if self.cached_penc is not None and self.cached_penc.shape == tensor.shape:
+ if num_objects is None:
+ return self.cached_penc
+ else:
+ return self.cached_penc.unsqueeze(1)
+
+ self.cached_penc = None
+
+ pos_y = torch.arange(h, device=tensor.device, dtype=self.inv_freq.dtype)
+ pos_x = torch.arange(w, device=tensor.device, dtype=self.inv_freq.dtype)
+ if self.normalize:
+ pos_y = pos_y / (pos_y[-1] + self.eps) * self.scale
+ pos_x = pos_x / (pos_x[-1] + self.eps) * self.scale
+
+ sin_inp_y = torch.einsum("i,j->ij", pos_y, self.inv_freq)
+ sin_inp_x = torch.einsum("i,j->ij", pos_x, self.inv_freq)
+ emb_y = get_emb(sin_inp_y).unsqueeze(1)
+ emb_x = get_emb(sin_inp_x)
+
+ emb = torch.zeros((h, w, self.dim * 2), device=tensor.device, dtype=tensor.dtype)
+ emb[:, :, :self.dim] = emb_x
+ emb[:, :, self.dim:] = emb_y
+
+ if not self.channel_last and self.transpose_output:
+ # cancelled out
+ pass
+ elif (not self.channel_last) or (self.transpose_output):
+ emb = emb.permute(2, 0, 1)
+
+ self.cached_penc = emb.unsqueeze(0).repeat(batch_size, 1, 1, 1)
+ if num_objects is None:
+ return self.cached_penc
+ else:
+ return self.cached_penc.unsqueeze(1)
+
+
+if __name__ == '__main__':
+ pe = PositionalEncoding(8).cuda()
+ input = torch.ones((1, 8, 8, 8)).cuda()
+ output = pe(input)
+ # print(output)
+ print(output[0, :, 0, 0])
+ print(output[0, :, 0, 5])
+ print(output[0, 0, :, 0])
+ print(output[0, 0, 0, :])
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/transformer_layers.py b/auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/transformer_layers.py
new file mode 100644
index 0000000..381f182
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/model/transformer/transformer_layers.py
@@ -0,0 +1,161 @@
+# Modified from PyTorch nn.Transformer
+
+from typing import List, Callable
+
+import torch
+from torch import Tensor
+import torch.nn as nn
+import torch.nn.functional as F
+from auto_media_publisher.modifiers.video_matting.matanyone.model.channel_attn import CAResBlock
+
+
+class SelfAttention(nn.Module):
+ def __init__(self,
+ dim: int,
+ nhead: int,
+ dropout: float = 0.0,
+ batch_first: bool = True,
+ add_pe_to_qkv: List[bool] = [True, True, False]):
+ super().__init__()
+ self.self_attn = nn.MultiheadAttention(dim, nhead, dropout=dropout, batch_first=batch_first)
+ self.norm = nn.LayerNorm(dim)
+ self.dropout = nn.Dropout(dropout)
+ self.add_pe_to_qkv = add_pe_to_qkv
+
+ def forward(self,
+ x: torch.Tensor,
+ pe: torch.Tensor,
+ attn_mask: bool = None,
+ key_padding_mask: bool = None) -> torch.Tensor:
+ x = self.norm(x)
+ if any(self.add_pe_to_qkv):
+ x_with_pe = x + pe
+ q = x_with_pe if self.add_pe_to_qkv[0] else x
+ k = x_with_pe if self.add_pe_to_qkv[1] else x
+ v = x_with_pe if self.add_pe_to_qkv[2] else x
+ else:
+ q = k = v = x
+
+ r = x
+ x = self.self_attn(q, k, v, attn_mask=attn_mask, key_padding_mask=key_padding_mask)[0]
+ return r + self.dropout(x)
+
+
+# https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html#torch.nn.functional.scaled_dot_product_attention
+class CrossAttention(nn.Module):
+ def __init__(self,
+ dim: int,
+ nhead: int,
+ dropout: float = 0.0,
+ batch_first: bool = True,
+ add_pe_to_qkv: List[bool] = [True, True, False],
+ residual: bool = True,
+ norm: bool = True):
+ super().__init__()
+ self.cross_attn = nn.MultiheadAttention(dim,
+ nhead,
+ dropout=dropout,
+ batch_first=batch_first)
+ if norm:
+ self.norm = nn.LayerNorm(dim)
+ else:
+ self.norm = nn.Identity()
+ self.dropout = nn.Dropout(dropout)
+ self.add_pe_to_qkv = add_pe_to_qkv
+ self.residual = residual
+
+ def forward(self,
+ x: torch.Tensor,
+ mem: torch.Tensor,
+ x_pe: torch.Tensor,
+ mem_pe: torch.Tensor,
+ attn_mask: bool = None,
+ *,
+ need_weights: bool = False) -> (torch.Tensor, torch.Tensor):
+ x = self.norm(x)
+ if self.add_pe_to_qkv[0]:
+ q = x + x_pe
+ else:
+ q = x
+
+ if any(self.add_pe_to_qkv[1:]):
+ mem_with_pe = mem + mem_pe
+ k = mem_with_pe if self.add_pe_to_qkv[1] else mem
+ v = mem_with_pe if self.add_pe_to_qkv[2] else mem
+ else:
+ k = v = mem
+ r = x
+ x, weights = self.cross_attn(q,
+ k,
+ v,
+ attn_mask=attn_mask,
+ need_weights=need_weights,
+ average_attn_weights=False)
+
+ if self.residual:
+ return r + self.dropout(x), weights
+ else:
+ return self.dropout(x), weights
+
+
+class FFN(nn.Module):
+ def __init__(self, dim_in: int, dim_ff: int, activation=F.relu):
+ super().__init__()
+ self.linear1 = nn.Linear(dim_in, dim_ff)
+ self.linear2 = nn.Linear(dim_ff, dim_in)
+ self.norm = nn.LayerNorm(dim_in)
+
+ if isinstance(activation, str):
+ self.activation = _get_activation_fn(activation)
+ else:
+ self.activation = activation
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ r = x
+ x = self.norm(x)
+ x = self.linear2(self.activation(self.linear1(x)))
+ x = r + x
+ return x
+
+
+class PixelFFN(nn.Module):
+ def __init__(self, dim: int):
+ super().__init__()
+ self.dim = dim
+ self.conv = CAResBlock(dim, dim)
+
+ def forward(self, pixel: torch.Tensor, pixel_flat: torch.Tensor) -> torch.Tensor:
+ # pixel: batch_size * num_objects * dim * H * W
+ # pixel_flat: (batch_size*num_objects) * (H*W) * dim
+ bs, num_objects, _, h, w = pixel.shape
+ pixel_flat = pixel_flat.view(bs * num_objects, h, w, self.dim)
+ pixel_flat = pixel_flat.permute(0, 3, 1, 2).contiguous()
+
+ x = self.conv(pixel_flat)
+ x = x.view(bs, num_objects, self.dim, h, w)
+ return x
+
+
+class OutputFFN(nn.Module):
+ def __init__(self, dim_in: int, dim_out: int, activation=F.relu):
+ super().__init__()
+ self.linear1 = nn.Linear(dim_in, dim_out)
+ self.linear2 = nn.Linear(dim_out, dim_out)
+
+ if isinstance(activation, str):
+ self.activation = _get_activation_fn(activation)
+ else:
+ self.activation = activation
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.linear2(self.activation(self.linear1(x)))
+ return x
+
+
+def _get_activation_fn(activation: str) -> Callable[[Tensor], Tensor]:
+ if activation == "relu":
+ return F.relu
+ elif activation == "gelu":
+ return F.gelu
+
+ raise RuntimeError("activation should be relu/gelu, not {}".format(activation))
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/__init__.py b/auto_media_publisher/modifiers/video_matting/matanyone/model/utils/__init__.py
similarity index 100%
copy from auto_media_publisher/modifiers/video_matting/matting_model/__init__.py
copy to auto_media_publisher/modifiers/video_matting/matanyone/model/utils/__init__.py
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/model/utils/memory_utils.py b/auto_media_publisher/modifiers/video_matting/matanyone/model/utils/memory_utils.py
new file mode 100644
index 0000000..e7dd5e7
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/model/utils/memory_utils.py
@@ -0,0 +1,107 @@
+import math
+import torch
+from typing import Optional, Union, Tuple
+
+
+# @torch.jit.script
+def get_similarity(mk: torch.Tensor,
+ ms: torch.Tensor,
+ qk: torch.Tensor,
+ qe: torch.Tensor,
+ add_batch_dim: bool = False,
+ uncert_mask = None) -> torch.Tensor:
+ # used for training/inference and memory reading/memory potentiation
+ # mk: B x CK x [N] - Memory keys
+ # ms: B x 1 x [N] - Memory shrinkage
+ # qk: B x CK x [HW/P] - Query keys
+ # qe: B x CK x [HW/P] - Query selection
+ # Dimensions in [] are flattened
+ # Return: B*N*HW
+ if add_batch_dim:
+ mk, ms = mk.unsqueeze(0), ms.unsqueeze(0)
+ qk, qe = qk.unsqueeze(0), qe.unsqueeze(0)
+
+ CK = mk.shape[1]
+
+ mk = mk.flatten(start_dim=2)
+ ms = ms.flatten(start_dim=1).unsqueeze(2) if ms is not None else None
+ qk = qk.flatten(start_dim=2)
+ qe = qe.flatten(start_dim=2) if qe is not None else None
+
+ # query token selection based on temporal sparsity
+ if uncert_mask is not None:
+ uncert_mask = uncert_mask.flatten(start_dim=2)
+ uncert_mask = uncert_mask.expand(-1, 64, -1)
+ qk = qk * uncert_mask
+ qe = qe * uncert_mask
+
+ if qe is not None:
+ # See XMem's appendix for derivation
+ mk = mk.transpose(1, 2)
+ a_sq = (mk.pow(2) @ qe)
+ two_ab = 2 * (mk @ (qk * qe))
+ b_sq = (qe * qk.pow(2)).sum(1, keepdim=True)
+ similarity = (-a_sq + two_ab - b_sq)
+ else:
+ # similar to STCN if we don't have the selection term
+ a_sq = mk.pow(2).sum(1).unsqueeze(2)
+ two_ab = 2 * (mk.transpose(1, 2) @ qk)
+ similarity = (-a_sq + two_ab)
+
+ if ms is not None:
+ similarity = similarity * ms / math.sqrt(CK) # B*N*HW
+ else:
+ similarity = similarity / math.sqrt(CK) # B*N*HW
+
+ return similarity
+
+
+def do_softmax(
+ similarity: torch.Tensor,
+ top_k: Optional[int] = None,
+ inplace: bool = False,
+ return_usage: bool = False) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+ # normalize similarity with top-k softmax
+ # similarity: B x N x [HW/P]
+ # use inplace with care
+ if top_k is not None:
+ values, indices = torch.topk(similarity, k=top_k, dim=1)
+
+ x_exp = values.exp_()
+ x_exp /= torch.sum(x_exp, dim=1, keepdim=True)
+ if inplace:
+ similarity.zero_().scatter_(1, indices, x_exp) # B*N*HW
+ affinity = similarity
+ else:
+ affinity = torch.zeros_like(similarity).scatter_(1, indices, x_exp) # B*N*HW
+ else:
+ maxes = torch.max(similarity, dim=1, keepdim=True)[0]
+ x_exp = torch.exp(similarity - maxes)
+ x_exp_sum = torch.sum(x_exp, dim=1, keepdim=True)
+ affinity = x_exp / x_exp_sum
+ indices = None
+
+ if return_usage:
+ return affinity, affinity.sum(dim=2)
+
+ return affinity
+
+
+def get_affinity(mk: torch.Tensor, ms: torch.Tensor, qk: torch.Tensor,
+ qe: torch.Tensor, uncert_mask = None) -> torch.Tensor:
+ # shorthand used in training with no top-k
+ similarity = get_similarity(mk, ms, qk, qe, uncert_mask=uncert_mask)
+ affinity = do_softmax(similarity)
+ return affinity
+
+def readout(affinity: torch.Tensor, mv: torch.Tensor, uncert_mask: torch.Tensor=None) -> torch.Tensor:
+ B, CV, T, H, W = mv.shape
+
+ mo = mv.view(B, CV, T * H * W)
+ mem = torch.bmm(mo, affinity)
+ if uncert_mask is not None:
+ uncert_mask = uncert_mask.flatten(start_dim=2).expand(-1, CV, -1)
+ mem = mem * uncert_mask
+ mem = mem.view(B, CV, H, W)
+
+ return mem
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/model/utils/parameter_groups.py b/auto_media_publisher/modifiers/video_matting/matanyone/model/utils/parameter_groups.py
new file mode 100644
index 0000000..177866a
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/model/utils/parameter_groups.py
@@ -0,0 +1,72 @@
+import logging
+
+log = logging.getLogger()
+
+
+def get_parameter_groups(model, stage_cfg, print_log=False):
+ """
+ Assign different weight decays and learning rates to different parameters.
+ Returns a parameter group which can be passed to the optimizer.
+ """
+ weight_decay = stage_cfg.weight_decay
+ embed_weight_decay = stage_cfg.embed_weight_decay
+ backbone_lr_ratio = stage_cfg.backbone_lr_ratio
+ base_lr = stage_cfg.learning_rate
+
+ backbone_params = []
+ embed_params = []
+ other_params = []
+
+ embedding_names = ['summary_pos', 'query_init', 'query_emb', 'obj_pe']
+ embedding_names = [e + '.weight' for e in embedding_names]
+
+ # inspired by detectron2
+ memo = set()
+ for name, param in model.named_parameters():
+ if not param.requires_grad:
+ continue
+ # Avoid duplicating parameters
+ if param in memo:
+ continue
+ memo.add(param)
+
+ if name.startswith('module'):
+ name = name[7:]
+
+ inserted = False
+ if name.startswith('pixel_encoder.'):
+ backbone_params.append(param)
+ inserted = True
+ if print_log:
+ log.info(f'{name} counted as a backbone parameter.')
+ else:
+ for e in embedding_names:
+ if name.endswith(e):
+ embed_params.append(param)
+ inserted = True
+ if print_log:
+ log.info(f'{name} counted as an embedding parameter.')
+ break
+
+ if not inserted:
+ other_params.append(param)
+
+ parameter_groups = [
+ {
+ 'params': backbone_params,
+ 'lr': base_lr * backbone_lr_ratio,
+ 'weight_decay': weight_decay
+ },
+ {
+ 'params': embed_params,
+ 'lr': base_lr,
+ 'weight_decay': embed_weight_decay
+ },
+ {
+ 'params': other_params,
+ 'lr': base_lr,
+ 'weight_decay': weight_decay
+ },
+ ]
+
+ return parameter_groups
\ No newline at end of file
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/model/utils/resnet.py b/auto_media_publisher/modifiers/video_matting/matanyone/model/utils/resnet.py
new file mode 100644
index 0000000..44886ee
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/model/utils/resnet.py
@@ -0,0 +1,179 @@
+"""
+resnet.py - A modified ResNet structure
+We append extra channels to the first conv by some network surgery
+"""
+
+from collections import OrderedDict
+import math
+
+import torch
+import torch.nn as nn
+from torch.utils import model_zoo
+
+
+def load_weights_add_extra_dim(target, source_state, extra_dim=1):
+ new_dict = OrderedDict()
+
+ for k1, v1 in target.state_dict().items():
+ if 'num_batches_tracked' not in k1:
+ if k1 in source_state:
+ tar_v = source_state[k1]
+
+ if v1.shape != tar_v.shape:
+ # Init the new segmentation channel with zeros
+ # print(v1.shape, tar_v.shape)
+ c, _, w, h = v1.shape
+ pads = torch.zeros((c, extra_dim, w, h), device=tar_v.device)
+ nn.init.orthogonal_(pads)
+ tar_v = torch.cat([tar_v, pads], 1)
+
+ new_dict[k1] = tar_v
+
+ target.load_state_dict(new_dict)
+
+
+model_urls = {
+ 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+ 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1, dilation=1):
+ return nn.Conv2d(in_planes,
+ out_planes,
+ kernel_size=3,
+ stride=stride,
+ padding=dilation,
+ dilation=dilation,
+ bias=False)
+
+
+class BasicBlock(nn.Module):
+ expansion = 1
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1):
+ super(BasicBlock, self).__init__()
+ self.conv1 = conv3x3(inplanes, planes, stride=stride, dilation=dilation)
+ self.bn1 = nn.BatchNorm2d(planes)
+ self.relu = nn.ReLU(inplace=True)
+ self.conv2 = conv3x3(planes, planes, stride=1, dilation=dilation)
+ self.bn2 = nn.BatchNorm2d(planes)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x):
+ residual = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+
+ if self.downsample is not None:
+ residual = self.downsample(x)
+
+ out += residual
+ out = self.relu(out)
+
+ return out
+
+
+class Bottleneck(nn.Module):
+ expansion = 4
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1):
+ super(Bottleneck, self).__init__()
+ self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(planes)
+ self.conv2 = nn.Conv2d(planes,
+ planes,
+ kernel_size=3,
+ stride=stride,
+ dilation=dilation,
+ padding=dilation,
+ bias=False)
+ self.bn2 = nn.BatchNorm2d(planes)
+ self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+ self.bn3 = nn.BatchNorm2d(planes * 4)
+ self.relu = nn.ReLU(inplace=True)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x):
+ residual = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+ out = self.relu(out)
+
+ out = self.conv3(out)
+ out = self.bn3(out)
+
+ if self.downsample is not None:
+ residual = self.downsample(x)
+
+ out += residual
+ out = self.relu(out)
+
+ return out
+
+
+class ResNet(nn.Module):
+ def __init__(self, block, layers=(3, 4, 23, 3), extra_dim=0):
+ self.inplanes = 64
+ super(ResNet, self).__init__()
+ self.conv1 = nn.Conv2d(3 + extra_dim, 64, kernel_size=7, stride=2, padding=3, bias=False)
+ self.bn1 = nn.BatchNorm2d(64)
+ self.relu = nn.ReLU(inplace=True)
+ self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+ self.layer1 = self._make_layer(block, 64, layers[0])
+ self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+ self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+ self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+
+ for m in self.modules():
+ if isinstance(m, nn.Conv2d):
+ n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+ m.weight.data.normal_(0, math.sqrt(2. / n))
+ elif isinstance(m, nn.BatchNorm2d):
+ m.weight.data.fill_(1)
+ m.bias.data.zero_()
+
+ def _make_layer(self, block, planes, blocks, stride=1, dilation=1):
+ downsample = None
+ if stride != 1 or self.inplanes != planes * block.expansion:
+ downsample = nn.Sequential(
+ nn.Conv2d(self.inplanes,
+ planes * block.expansion,
+ kernel_size=1,
+ stride=stride,
+ bias=False),
+ nn.BatchNorm2d(planes * block.expansion),
+ )
+
+ layers = [block(self.inplanes, planes, stride, downsample)]
+ self.inplanes = planes * block.expansion
+ for i in range(1, blocks):
+ layers.append(block(self.inplanes, planes, dilation=dilation))
+
+ return nn.Sequential(*layers)
+
+
+def resnet18(pretrained=True, extra_dim=0):
+ model = ResNet(BasicBlock, [2, 2, 2, 2], extra_dim)
+ if pretrained:
+ load_weights_add_extra_dim(model, model_zoo.load_url(model_urls['resnet18']), extra_dim)
+ return model
+
+
+def resnet50(pretrained=True, extra_dim=0):
+ model = ResNet(Bottleneck, [3, 4, 6, 3], extra_dim)
+ if pretrained:
+ load_weights_add_extra_dim(model, model_zoo.load_url(model_urls['resnet50']), extra_dim)
+ return model
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/__init__.py b/auto_media_publisher/modifiers/video_matting/matanyone/utils/__init__.py
similarity index 100%
copy from auto_media_publisher/modifiers/video_matting/matting_model/__init__.py
copy to auto_media_publisher/modifiers/video_matting/matanyone/utils/__init__.py
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/utils/get_default_model.py b/auto_media_publisher/modifiers/video_matting/matanyone/utils/get_default_model.py
new file mode 100644
index 0000000..1f67564
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/utils/get_default_model.py
@@ -0,0 +1,27 @@
+"""
+A helper function to get a default model for quick testing
+"""
+from omegaconf import open_dict
+from hydra import compose, initialize
+
+import torch
+from auto_media_publisher.modifiers.video_matting.matanyone.model.matanyone import MatAnyone
+
+def get_matanyone_model(ckpt_path, device=None) -> MatAnyone:
+ initialize(version_base='1.3.2', config_path="../config", job_name="eval_our_config")
+ cfg = compose(config_name="eval_matanyone_config")
+
+ with open_dict(cfg):
+ cfg['weights'] = ckpt_path
+
+ # Load the network weights
+ if device is not None:
+ matanyone = MatAnyone(cfg, single_object=True).to(device).eval()
+ model_weights = torch.load(cfg.weights, map_location=device)
+ else: # if device is not specified, `.cuda()` by default
+ matanyone = MatAnyone(cfg, single_object=True).cuda().eval()
+ model_weights = torch.load(cfg.weights)
+
+ matanyone.load_weights(model_weights)
+
+ return matanyone
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/utils/inference_utils.py b/auto_media_publisher/modifiers/video_matting/matanyone/utils/inference_utils.py
new file mode 100644
index 0000000..7bfcb77
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/utils/inference_utils.py
@@ -0,0 +1,57 @@
+import os
+import cv2
+import random
+import numpy as np
+
+import torch
+import torchvision
+
+IMAGE_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG')
+VIDEO_EXTENSIONS = ('.mp4', '.mov', '.avi', '.MP4', '.MOV', '.AVI')
+
+def read_frame_from_videos(frame_root):
+ if frame_root.endswith(VIDEO_EXTENSIONS): # Video file path
+ video_name = os.path.basename(frame_root)[:-4]
+ frames, _, info = torchvision.io.read_video(filename=frame_root, pts_unit='sec', output_format='TCHW') # RGB
+ fps = info['video_fps']
+ height, width = frames.shape[2], frames.shape[3]
+ else:
+ video_name = os.path.basename(frame_root)
+ frames = []
+ fr_lst = sorted(os.listdir(frame_root))
+ for fr in fr_lst:
+ frame = cv2.imread(os.path.join(frame_root, fr))[...,[2,1,0]] # RGB, HWC
+ frames.append(frame)
+ fps = 30 # default
+ height, width = frames.shape[2], frames.shape[3]
+ frames = torch.Tensor(np.array(frames)).permute(0, 3, 1, 2).contiguous() # TCHW
+
+ length = frames.shape[0]
+
+ return frames, fps,width,height,length, video_name
+
+
+def get_video_paths(input_root):
+ video_paths = []
+ for root, _, files in os.walk(input_root):
+ for file in files:
+ if file.lower().endswith(VIDEO_EXTENSIONS):
+ video_paths.append(os.path.join(root, file))
+ return sorted(video_paths)
+
+def str_to_list(value):
+ return list(map(int, value.split(',')))
+
+def gen_dilate(alpha, min_kernel_size, max_kernel_size):
+ kernel_size = random.randint(min_kernel_size, max_kernel_size)
+ kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size,kernel_size))
+ fg_and_unknown = np.array(np.not_equal(alpha, 0).astype(np.float32))
+ dilate = cv2.dilate(fg_and_unknown, kernel, iterations=1)*255
+ return dilate.astype(np.float32)
+
+def gen_erosion(alpha, min_kernel_size, max_kernel_size):
+ kernel_size = random.randint(min_kernel_size, max_kernel_size)
+ kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size,kernel_size))
+ fg = np.array(np.equal(alpha, 255).astype(np.float32))
+ erode = cv2.erode(fg, kernel, iterations=1)*255
+ return erode.astype(np.float32)
\ No newline at end of file
diff --git a/auto_media_publisher/modifiers/video_matting/matanyone/utils/tensor_utils.py b/auto_media_publisher/modifiers/video_matting/matanyone/utils/tensor_utils.py
new file mode 100644
index 0000000..52e8a9a
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matanyone/utils/tensor_utils.py
@@ -0,0 +1,62 @@
+from typing import List, Iterable
+import torch
+import torch.nn.functional as F
+
+
+# STM
+def pad_divide_by(in_img: torch.Tensor, d: int) -> (torch.Tensor, Iterable[int]):
+ h, w = in_img.shape[-2:]
+
+ if h % d > 0:
+ new_h = h + d - h % d
+ else:
+ new_h = h
+ if w % d > 0:
+ new_w = w + d - w % d
+ else:
+ new_w = w
+ lh, uh = int((new_h - h) / 2), int(new_h - h) - int((new_h - h) / 2)
+ lw, uw = int((new_w - w) / 2), int(new_w - w) - int((new_w - w) / 2)
+ pad_array = (int(lw), int(uw), int(lh), int(uh))
+ out = F.pad(in_img, pad_array)
+ return out, pad_array
+
+
+def unpad(img: torch.Tensor, pad: Iterable[int]) -> torch.Tensor:
+ if len(img.shape) == 4:
+ if pad[2] + pad[3] > 0:
+ img = img[:, :, pad[2]:-pad[3], :]
+ if pad[0] + pad[1] > 0:
+ img = img[:, :, :, pad[0]:-pad[1]]
+ elif len(img.shape) == 3:
+ if pad[2] + pad[3] > 0:
+ img = img[:, pad[2]:-pad[3], :]
+ if pad[0] + pad[1] > 0:
+ img = img[:, :, pad[0]:-pad[1]]
+ elif len(img.shape) == 5:
+ if pad[2] + pad[3] > 0:
+ img = img[:, :, :, pad[2]:-pad[3], :]
+ if pad[0] + pad[1] > 0:
+ img = img[:, :, :, :, pad[0]:-pad[1]]
+ else:
+ raise NotImplementedError
+ return img
+
+
+# @torch.jit.script
+def aggregate(prob: torch.Tensor, dim: int) -> torch.Tensor:
+ with torch.amp.autocast("cuda",enabled=False):
+ prob = prob.float()
+ new_prob = torch.cat([torch.prod(1 - prob, dim=dim, keepdim=True), prob],
+ dim).clamp(1e-7, 1 - 1e-7)
+ logits = torch.log((new_prob / (1 - new_prob))) # (0, 1) --> (-inf, inf)
+
+ return logits
+
+
+# @torch.jit.script
+def cls_to_one_hot(cls_gt: torch.Tensor, num_objects: int) -> torch.Tensor:
+ # cls_gt: B*1*H*W
+ B, _, H, W = cls_gt.shape
+ one_hot = torch.zeros(B, num_objects + 1, H, W, device=cls_gt.device).scatter_(1, cls_gt, 1)
+ return one_hot
\ No newline at end of file
diff --git a/auto_media_publisher/modifiers/video_matting/video_matting_converter.py b/auto_media_publisher/modifiers/video_matting/robust_video_matting_converter.py
similarity index 92%
rename from auto_media_publisher/modifiers/video_matting/video_matting_converter.py
rename to auto_media_publisher/modifiers/video_matting/robust_video_matting_converter.py
index f8662d6..85b69ca 100644
--- a/auto_media_publisher/modifiers/video_matting/video_matting_converter.py
+++ b/auto_media_publisher/modifiers/video_matting/robust_video_matting_converter.py
@@ -7,7 +7,7 @@
import time
from pathlib import Path
from auto_media_publisher.utils.media_input_output import VideoReader, VideoWriter, ImageSequenceReader, ImageSequenceWriter
-from .matting_model.model import MattingNetwork
+from .robustvideomatting.model import MattingNetwork
import uuid
@@ -18,7 +18,7 @@
CHECKPOINT_MOBILENETV3 ="./models/rvm_mobilenetv3.pth"
-class VideoMattingConverter:
+class RobustVideoMattingConverter:
"""
瑙嗛鎶犲儚杞崲鍣�
灏佽鍘熸潵鐨� convert_video 閫昏緫锛屾敮鎸� video / png_sequence 杈撳嚭
@@ -43,7 +43,7 @@
def video_matting(self,
input: str,
bg_video: Optional[str] = None,
- bg_jpg: Optional[str] = None,
+ bg_image: Optional[str] = None,
output_video: Optional[str] = None,
output_alpha: Optional[str] = None,
output_fore: Optional[str] = None):
@@ -52,10 +52,11 @@
if not os.path.exists(input):
print("input file is NOT exist.")
return False
+
if bg_video and not os.path.exists(bg_video):
print("input background video is NOT exist.")
return False
- if bg_jpg and not os.path.exists(bg_jpg):
+ if bg_image and not os.path.exists(bg_image):
print("input background image is NOT exist.")
return False
@@ -65,7 +66,7 @@
downsample_ratio = None,
output_type = 'video',
background_video = bg_video,
- background_jpg = bg_jpg,
+ background_jpg = bg_image,
with_audio = True,
output_composition = output_video,
output_alpha = output_alpha,
@@ -116,13 +117,13 @@
if output_composition:
if with_audio:
- writer_com = VideoWriter(output_composition, frame_rate, int(output_video_mbps * 1_000_000),input_source)
+ writer_com = VideoWriter(output_composition, frame_rate, source.width,source.height,int(output_video_mbps * 1_000_000),input_source)
else:
- writer_com = VideoWriter(output_composition, frame_rate, int(output_video_mbps * 1_000_000))
+ writer_com = VideoWriter(output_composition, frame_rate, source.width,source.height,int(output_video_mbps * 1_000_000))
if output_alpha:
- writer_pha = VideoWriter(output_alpha, frame_rate, int(output_video_mbps * 1_000_000))
+ writer_pha = VideoWriter(output_alpha, frame_rate, source.width,source.height,int(output_video_mbps * 1_000_000))
if output_foreground:
- writer_fgr = VideoWriter(output_foreground, frame_rate, int(output_video_mbps * 1_000_000))
+ writer_fgr = VideoWriter(output_foreground, frame_rate,source.width,source.height, int(output_video_mbps * 1_000_000))
else: # png_sequence
if output_composition:
@@ -141,7 +142,8 @@
bgr_IMG = Image.open(background_jpg)
bgr = transforms.ToTensor()(bgr_IMG).to(self.device)
else:
- bgr = torch.tensor([155, 255, 155], device=self.device, dtype=dtype).div(255).view(1, 1, 3, 1, 1)
+ # 浣跨敤 0,200,0 鑰屼笉鏄� 0,255,0 閬垮厤杩囧害鏇濆厜
+ bgr = torch.tensor([0, 200, 0], device=self.device, dtype=dtype).div(255).view(1, 1, 3, 1, 1)
# 寮�濮嬫帹鐞�
self.model = self.model.eval()
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/__init__.py b/auto_media_publisher/modifiers/video_matting/robustvideomatting/__init__.py
similarity index 100%
copy from auto_media_publisher/modifiers/video_matting/matting_model/__init__.py
copy to auto_media_publisher/modifiers/video_matting/robustvideomatting/__init__.py
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/decoder.py b/auto_media_publisher/modifiers/video_matting/robustvideomatting/decoder.py
similarity index 100%
rename from auto_media_publisher/modifiers/video_matting/matting_model/decoder.py
rename to auto_media_publisher/modifiers/video_matting/robustvideomatting/decoder.py
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/deep_guided_filter.py b/auto_media_publisher/modifiers/video_matting/robustvideomatting/deep_guided_filter.py
similarity index 100%
rename from auto_media_publisher/modifiers/video_matting/matting_model/deep_guided_filter.py
rename to auto_media_publisher/modifiers/video_matting/robustvideomatting/deep_guided_filter.py
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/fast_guided_filter.py b/auto_media_publisher/modifiers/video_matting/robustvideomatting/fast_guided_filter.py
similarity index 100%
rename from auto_media_publisher/modifiers/video_matting/matting_model/fast_guided_filter.py
rename to auto_media_publisher/modifiers/video_matting/robustvideomatting/fast_guided_filter.py
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/lraspp.py b/auto_media_publisher/modifiers/video_matting/robustvideomatting/lraspp.py
similarity index 100%
rename from auto_media_publisher/modifiers/video_matting/matting_model/lraspp.py
rename to auto_media_publisher/modifiers/video_matting/robustvideomatting/lraspp.py
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/mobilenetv3.py b/auto_media_publisher/modifiers/video_matting/robustvideomatting/mobilenetv3.py
similarity index 100%
rename from auto_media_publisher/modifiers/video_matting/matting_model/mobilenetv3.py
rename to auto_media_publisher/modifiers/video_matting/robustvideomatting/mobilenetv3.py
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/model.py b/auto_media_publisher/modifiers/video_matting/robustvideomatting/model.py
similarity index 100%
rename from auto_media_publisher/modifiers/video_matting/matting_model/model.py
rename to auto_media_publisher/modifiers/video_matting/robustvideomatting/model.py
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/resnet.py b/auto_media_publisher/modifiers/video_matting/robustvideomatting/resnet.py
similarity index 100%
rename from auto_media_publisher/modifiers/video_matting/matting_model/resnet.py
rename to auto_media_publisher/modifiers/video_matting/robustvideomatting/resnet.py
diff --git a/auto_media_publisher/utils/media_input_output.py b/auto_media_publisher/utils/media_input_output.py
index 0adf1bd..f4a9af6 100644
--- a/auto_media_publisher/utils/media_input_output.py
+++ b/auto_media_publisher/utils/media_input_output.py
@@ -5,12 +5,21 @@
from torch.utils.data import Dataset
from torchvision.transforms.functional import to_pil_image
from PIL import Image
-
+import imageio.v3 as iio
+import uuid
+import traceback
+import subprocess
+import time
class VideoReader(Dataset):
def __init__(self, path, transform=None):
self.video = pims.PyAVVideoReader(path)
self.rate = self.video.frame_rate
+
+ # 鐩存帴鎷跨涓�甯х殑 shape锛圚, W, C锛�
+ first_frame = np.asarray(self.video[0])
+ self.height, self.width = first_frame.shape[:2]
+
self.transform = transform
@property
@@ -29,23 +38,37 @@
class VideoWriter:
- def __init__(self, path, frame_rate, bit_rate=1000000,audio_source_path=None):
+ def __init__(self, output_path, frame_rate, width,height, bit_rate=1000000,audio_source_path=None):
"""
:param path: 杈撳嚭鏂囦欢璺緞
:param frame_rate: 瑙嗛甯х巼
:param bit_rate: 瑙嗛鐮佺巼
:param audio_source_path: 鍘熻棰戞枃浠惰矾寰勶紙鐢ㄤ簬鎷疯礉闊抽锛�
"""
- self.container = av.open(path, mode='w')
+ print(f"ready to create video file : {output_path}")
+ self.output_path = output_path
+ self.audio_source_path =audio_source_path
+ if self.audio_source_path and os.path.exists(self.audio_source_path):
+ # using temp video file without audio
+ self.need_audio_action = True
+ self.temp_path = f"{uuid.uuid4()}.mp4"
+ print(self.temp_path)
+ self.container = av.open(self.temp_path, mode='w')
+ else:
+ self.container = av.open(self.output_path, mode='w')
+ self.need_audio_action = False
+
self.v_stream = self.container.add_stream('h264', rate=round(frame_rate))
self.v_stream.pix_fmt = 'yuv420p'
self.v_stream.bit_rate = bit_rate
+ self.v_stream.width = width
+ self.v_stream.height = height
self.v_stream.codec_context.options = {
'crf': '18', # 鎺у埗鐢昏川锛岃秺灏忚秺娓呮櫚锛�18鏄珮璐ㄩ噺锛�
'preset': 'veryfast', # 褰卞搷缂栫爜閫熷害锛宮edium 鏄�1x锛宖ast鏄�2x faster鏄�3x veryfast鏄�4x superfast鏄�6x ultrafast鏄�10x锛屾枃浠朵綋绉ぇ绾︿細澧炲姞(浠edium涓哄熀鍑�) 10%锛�15%锛�20% 30% 50%
'tune': 'film',
}
- self.audio_source_path =audio_source_path
+ '''
self.audio_in_container = None
self.audio_stream_in = None
self.audio_out_stream = None
@@ -62,39 +85,134 @@
self.audio_out_stream = self.container.add_stream(self.audio_stream_in.codec.name)
else:
self.audio_out_stream = None
+ '''
-
+
def write(self, frames):
- # frames: [T, C, H, W]
- self.v_stream.width = frames.size(3)
- self.v_stream.height = frames.size(2)
+
if frames.size(1) == 1:
- frames = frames.repeat(1, 3, 1, 1) # convert grayscale to RGB
+ frames = frames.repeat(1, 3, 1, 1)
+
frames = frames.mul(255).byte().cpu().permute(0, 2, 3, 1).numpy()
+
for t in range(frames.shape[0]):
frame = frames[t]
- frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
- # old version
- #self.container.mux(self.v_stream.encode(frame))
- for packet in self.v_stream.encode(frame):
+ video_frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
+ for packet in self.v_stream.encode(video_frame):
self.container.mux(packet)
-
- def close(self):
- # old version
- #self.container.mux(self.v_stream.encode())
- for packet in self.v_stream.encode():
- self.container.mux(packet)
+
+ def write_nparray(self,frames):
+ try:
+ if frames.ndim == 3: # 鍗曞抚
+ frames = np.expand_dims(frames, axis=0)
+
+ if frames.shape[-1] == 1: # 鐏板害鍥�
+ frames = np.repeat(frames, 3, axis=-1) # HWC 鈫� RGB
+
+ for t in range(frames.shape[0]):
+ frame = frames[t]
+ # 鐩存帴鐢╪umpy鍐欏叆
+ video_frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
+ for packet in self.v_stream.encode(video_frame):
+ self.container.mux(packet)
+ except Exception as e:
+ print(f"{e}")
+ traceback.format_exc() # 鎵撳嵃鍫嗘爤淇℃伅锛堝鏋� logger 鏀寔 debug锛�
+
+
+
+ def merge_video_audio(self,video_path, audio_path, output_path):
+ """
+ 鐢� FFmpeg 鎶� video_path 鐨勭敾闈� 鍜� audio_path 鐨勯煶棰� 鍚堝苟鎴� output_path
+ 鏈�鐭殑涓�涓粨鏉� (-shortest)锛屽悎骞跺悗鍒犻櫎 video_path
+ """
+ try:
+ # 鏋勯�� FFmpeg 鍛戒护
+ cmd = [
+ "ffmpeg", "-y", # -y 瑕嗙洊杈撳嚭
+ "-loglevel", "error", # 鍙樉绀洪敊璇�
+ "-i", video_path, # 鐢婚潰
+ "-i", audio_path, # 闊抽
+ "-c:v", "copy", # 瑙嗛涓嶉噸鏂扮紪鐮�
+ "-c:a", "copy", # 闊抽涓嶉噸鏂扮紪鐮�
+ "-map", "0:v:0", # 绗竴涓緭鍏ョ殑瑙嗛
+ "-map", "1:a:0", # 绗簩涓緭鍏ョ殑闊抽
+ "-shortest", # 浠ユ渶鐭殑杈撳叆缁撴潫
+ output_path
+ ]
+
+ print(f"[INFO] 寮�濮嬪悎骞惰棰戝拰闊抽...\n{' '.join(cmd)}")
+ subprocess.run(cmd, check=True)
+
+ print("[INFO] 鍚堝苟瀹屾垚锛岀瓑寰� 1 绉掑悗鍒犻櫎鍘熺敾闈㈣棰�...")
+ time.sleep(1)
+
+ # 鍒犻櫎鍘熻棰�
+ if os.path.exists(video_path):
+ os.remove(video_path)
+ print(f"[INFO] 宸插垹闄ゅ師鐢婚潰瑙嗛: {video_path}")
+
+ except subprocess.CalledProcessError as e:
+ print(f"[ERROR] FFmpeg 鍚堝苟澶辫触锛岄敊璇爜: {e.returncode}")
+ except Exception as e:
+ print(f"[ERROR] 鍙戠敓鏈煡閿欒: {e}")
+
+
+
+ def close(self):
+ '''
if self.audio_stream_in and self.audio_out_stream:
for packet in self.audio_in_container.demux(self.audio_stream_in):
if packet.dts is None:
continue
packet.stream = self.audio_out_stream # 鐩存帴鏀癸紝涓嶇敤 copy()
self.container.mux(packet)
-
- self.container.close()
if self.audio_in_container:
self.audio_in_container.close()
+ '''
+ self.container.close()
+
+ if self.need_audio_action:
+ self.merge_video_audio(self.temp_path,self.audio_source_path,self.output_path)
+
+
+
+class WebMWriter:
+ """
+ 鏀寔閫忔槑閫氶亾鐨勪腑闂磋棰戞牸寮�
+ # 绀轰緥浣跨敤锛�
+ writer = WebMWriter("output_alpha.webm", fps=30)
+ for i in range(frames):
+ writer.write(frame)
+ writer.close()
+ """
+ def __init__(self, filename, fps=30):
+ self.filename = filename
+ self.fps = fps
+ self.frames = []
+
+ def write(self, frame: np.ndarray):
+ """
+ 浼犲叆鐨� frame 搴旇鏄惈 alpha 閫氶亾鐨� numpy 鏁扮粍锛宻hape=(H,W,4)锛宒type=uint8
+ RGBA 鏍煎紡锛宎lpha 閫氶亾鍐冲畾閫忔槑搴�
+ """
+ assert frame.shape[2] == 4, "Frame must have 4 channels (RGBA)"
+ self.frames.append(frame)
+
+ def close(self):
+ iio.imwrite(
+ self.filename,
+ self.frames,
+ fps=self.fps,
+ codec="vp8", # vp8鏀寔alpha閫氶亾
+ ffmpeg_params=[
+ "-auto-alt-ref", "0", # 鍏抽棴鍙傝�冨抚锛屾敮鎸侀�忔槑
+ "-pix_fmt", "yuva420p" # VP8閫忔槑鏍煎紡
+ ],
+ quality=10 # 璐ㄩ噺锛岃寖鍥�1-63锛�10姣旇緝楂樿川閲�
+ )
+ print(f"WebM video saved to {self.filename}")
class ImageSequenceReader(Dataset):
diff --git a/auto_media_publisher/utils/utils_images.py b/auto_media_publisher/utils/utils_images.py
new file mode 100644
index 0000000..a1d2e28
--- /dev/null
+++ b/auto_media_publisher/utils/utils_images.py
@@ -0,0 +1,56 @@
+import io
+import numpy as np
+from PIL import Image
+
+
+def load_image_any(source):
+ """缁熶竴鍔犺浇涓� Pillow Image 瀵硅薄锛屽苟杩斿洖鍘熺被鍨嬫爣璁�"""
+ if isinstance(source, Image.Image):
+ return source.convert("RGBA"), "image"
+
+ if isinstance(source, np.ndarray):
+ return Image.fromarray(source).convert("RGBA"), "ndarray"
+
+ if isinstance(source, bytes):
+ return Image.open(io.BytesIO(source)).convert("RGBA"), "bytes"
+
+ raise ValueError("涓嶆敮鎸佺殑鍥剧墖杈撳叆绫诲瀷")
+
+def resize_and_crop_center_any(source, target_w, target_h):
+ img, src_type = load_image_any(source)
+ if img is None:
+ return None
+
+ src_w, src_h = img.size
+ target_ratio = target_w / target_h
+ src_ratio = src_w / src_h
+
+ # 璁＄畻缂╂斁姣斾緥锛堜繚璇佷竴杈瑰垰濂斤紝鍙︿竴杈瑰鍑烘潵锛�
+ if src_ratio > target_ratio:
+ scale = target_h / src_h # 鎸夐珮搴﹀尮閰�
+ else:
+ scale = target_w / src_w # 鎸夊搴﹀尮閰�
+
+ new_w = int(src_w * scale)
+ new_h = int(src_h * scale)
+ img_resized = img.resize((new_w, new_h), Image.LANCZOS)
+
+ # 灞呬腑瑁佸壀
+ left = max((new_w - target_w) / 2, 0)
+ top = max((new_h - target_h) / 2, 0)
+ right = left + target_w
+ bottom = top + target_h
+ img_cropped = img_resized.crop((left, top, right, bottom))
+
+ # 杩斿洖涓庡師杈撳叆涓�鑷寸殑绫诲瀷
+ if src_type == "image":
+ return img_cropped
+ elif src_type == "bytes":
+ buf = io.BytesIO()
+ img_cropped.save(buf, format="PNG")
+ return buf.getvalue()
+ elif src_type == "ndarray":
+ return np.array(img_cropped)
+ else:
+ return img_cropped # 鍏滃簳锛岃繑鍥� Pillow Image
+
diff --git a/auto_media_publisher/utils/utils_text2image_pollinations.py b/auto_media_publisher/utils/utils_text2image_pollinations.py
index f78411b..6f107e5 100644
--- a/auto_media_publisher/utils/utils_text2image_pollinations.py
+++ b/auto_media_publisher/utils/utils_text2image_pollinations.py
@@ -12,6 +12,7 @@
POLLINATIONS_DEFAULT_HEIGHT ,
POLLINATIONS_DEFAULT_MODEL,
)
+from .utils_images import resize_and_crop_center_any
class Text2ImageGenerator:
@@ -61,12 +62,12 @@
try:
response = requests.get(url, timeout=30)
response.raise_for_status()
- return response.content
+ return resize_and_crop_center_any(response.content,self.width,self.height)
except Exception as e:
print(f"[!] 璇锋眰绗� {attempt} 娆″け璐ワ細{e}")
return None
- def generate_image(self, prompt: str, filename: str = None) -> str:
+ def generate_image_with_trans(self, prompt: str, filename: str = None) -> str:
if not prompt.strip():
raise ValueError("鎻愮ず璇嶄笉鑳戒负绌�")
@@ -90,7 +91,31 @@
print(f"[鉁揮 鍥剧墖宸蹭繚瀛橈細{abs_path}")
return abs_path
- def generate_image_to_stream(self, prompt: str) -> BytesIO:
+ def generate_image_without_trans(self, prompt: str, filename: str = None) -> str:
+ if not prompt.strip():
+ raise ValueError("鎻愮ず璇嶄笉鑳戒负绌�")
+
+ url = self.build_url(prompt)
+ print(url)
+
+ content = self._fetch_with_retry(url)
+ if content is None:
+ return None
+
+ if not filename:
+ now = datetime.datetime.now()
+ filename = now.strftime("%Y%m%d%H%M%S") + ".png"
+ if not filename.lower().endswith(".png"):
+ filename += ".png"
+
+ with open(filename, "wb") as f:
+ f.write(content)
+
+ abs_path = os.path.abspath(filename)
+ print(f"[鉁揮 鍥剧墖宸蹭繚瀛橈細{abs_path}")
+ return abs_path
+
+ def generate_image_with_trans_to_stream(self, prompt: str) -> BytesIO:
if not prompt.strip():
raise ValueError("鎻愮ず璇嶄笉鑳戒负绌�")
diff --git a/auto_media_publisher/utils/utils_yolo.py b/auto_media_publisher/utils/utils_yolo.py
new file mode 100644
index 0000000..72bafeb
--- /dev/null
+++ b/auto_media_publisher/utils/utils_yolo.py
@@ -0,0 +1,79 @@
+import os
+import cv2
+import numpy as np
+from datetime import datetime
+from ultralytics import YOLO
+
+
+DEFAULT_YOLOV8S_SEG_MODEL = './models/yolov8s-seg.pt'
+
+def get_center_distance(box, image_center):
+ x1, y1, x2, y2 = box
+ box_center = ((x1 + x2) / 2, (y1 + y2) / 2)
+ return np.linalg.norm(np.array(box_center) - np.array(image_center))
+
+def generate_mask_from_first_frame(video_path, filter_type:str = 'dist',mask_type:str="array",mask_file_path:str=None):
+ """
+ 鑾峰彇瑙嗛绗竴甯х殑 浜哄儚 MASK:
+ filter_type锛�
+ dist 璺濈 鐢婚潰涓績鏈�杩戠殑浜�; default
+ box : 鐢婚潰涓潰绉渶澶х殑浜�;
+ mask_type :
+ array: 鐩存帴杩斿洖 np.array ; default
+ file锛氫繚瀛樻垚mask_file_path 锛岃繑鍥瀖ask_file_path
+ """
+ if not os.path.exists(video_path):
+ raise FileNotFoundError(f"瑙嗛鏂囦欢涓嶅瓨鍦細{video_path}")
+
+ cap = cv2.VideoCapture(video_path)
+ ret, frame = cap.read()
+ cap.release()
+
+ if not ret:
+ raise RuntimeError("鏃犳硶璇诲彇瑙嗛绗竴甯�")
+
+ height, width, _ = frame.shape
+ image_center = (width / 2, height / 2)
+
+
+ print(os.path.exists(DEFAULT_YOLOV8S_SEG_MODEL))
+
+ model = YOLO(DEFAULT_YOLOV8S_SEG_MODEL) # 鎴� yolov8s-seg.pt锛屼綋绉洿澶т絾鏇村噯
+
+ # 3. 杩涜浜哄舰妫�娴�+鍒嗗壊
+ results = model(frame)[0]
+
+ # 绛涢�夊嚭鎵�鏈夌被鍒负鈥減erson鈥濈殑鐩爣
+ person_masks = []
+ for i, cls in enumerate(results.boxes.cls):
+ if int(cls) == 0: # class 0 鏄� "person"
+ box = results.boxes.xyxy[i].cpu().numpy() # [x1, y1, x2, y2]
+ mask = results.masks.data[i].cpu().numpy()
+
+ dist = get_center_distance(box, image_center)
+ box_area = (box[2] - box[0]) * (box[3] - box[1])
+
+ person_masks.append((dist, box_area, mask))
+
+ if not person_masks:
+ return None
+
+ # 4. 鍙栬窛绂讳腑蹇冩渶杩戠殑浜哄舰
+ if filter_type=="box":
+ person_masks.sort(key=lambda x: x[1], reverse=True)
+ else:
+ person_masks.sort(key=lambda x: x[0])
+
+ best_mask = person_masks[0][2]
+ best_mask_resized = cv2.resize(best_mask, (width, height), interpolation=cv2.INTER_NEAREST)
+
+ if mask_type == "file":
+ if not mask_file_path:
+ now_str = datetime.now().strftime("%Y%m%d_%H%M%S")
+ mask_file_path =f"./temp/{now_str}.png"
+ mask_img = (best_mask_resized * 255).astype(np.uint8)
+ cv2.imwrite(mask_file_path, mask_img)
+ return mask_file_path
+ else:
+ mask = (best_mask_resized * 255).astype(np.uint8)
+ return mask
\ No newline at end of file
diff --git a/data/auto_media_publisher.sqlite3 b/data/auto_media_publisher.sqlite3
index 1364c4f..70ca126 100644
--- a/data/auto_media_publisher.sqlite3
+++ b/data/auto_media_publisher.sqlite3
Binary files differ
diff --git a/data/cookies/18918594752/gongzhonghao.json b/data/cookies/18918594752/gongzhonghao.json
index 588d31f..50477ae 100644
--- a/data/cookies/18918594752/gongzhonghao.json
+++ b/data/cookies/18918594752/gongzhonghao.json
@@ -1 +1 @@
-{"cookies": [{"name": "ua_id", "value": "Yp8GoVBn7Tj0PWR7AAAAADaCnYPYpvql9WPFzFkbL_A=", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1788489057.928054, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "wxuin", "value": "53845090997782", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1788405090.13786, "httpOnly": false, "secure": false, "sameSite": "Lax"}, {"name": "uuid", "value": "8093d628ddf047d5a340ec437ca7fd08", "domain": "mp.weixin.qq.com", "path": "/", "expires": -1, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "rand_info", "value": "CAESIEjp7aANF1MLCho4lIEi6Ju8oYLEexqCUxELmbqpfG3Y", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754274687, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "slave_bizuin", "value": "3932822909", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754274687, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "data_bizuin", "value": "3932822909", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754274687, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "bizuin", "value": "3932822909", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754274687, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "data_ticket", "value": "TXXMQ0chAAhF8D9olYk9AhfbDAjnS2jHUUO8GHs0e45nt6P/zjA38GlY5q12SD0E", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754274687, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "slave_sid", "value": "OHl1NHRHZjd2N3poQTFDQ2d5WVBBWnJ5bm9obkZNVjZGank5S2dxSUMwblRqeFUwalBSNEpTbHlrMjZ6cUtxT2dRWk9NdElDcGRmUko1a1F4MENlUWFja0FvVjZDZnY0RXVkT0NiNUJnOVlITGQzOHRGYTBLUG10UEZzSzJhdWozakN1SUo5VXZRNmdkUGlz", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754274687, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "slave_user", "value": "gh_5624b93dc545", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754274687, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "xid", "value": "fedc4a4eeb3b124f5812d5b88b1265f4", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1788405116.292443, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "mm_lang", "value": "zh_CN", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1788405116.292574, "httpOnly": false, "secure": true, "sameSite": "Lax"}, {"name": "_clck", "value": "3932822909|1|fy2|0", "domain": ".qq.com", "path": "/", "expires": 1785464417, "httpOnly": false, "secure": false, "sameSite": "Lax"}, {"name": "_clsk", "value": "184wq4q|1753929064020|7|1|mp.weixin.qq.com/weheat-agent/payload/record", "domain": ".qq.com", "path": "/", "expires": 1754015464, "httpOnly": false, "secure": false, "sameSite": "Lax"}], "origins": [{"origin": "https://mp.weixin.qq.com", "localStorage": [{"name": "USER_AGENT_BROWSER_MAJOR_VERSION", "value": "127"}, {"name": "loginMode", "value": "1"}, {"name": "adTransitionText", "value": "noAdText"}, {"name": "hasClick", "value": "[0]"}]}]}
\ No newline at end of file
+{"cookies": [{"name": "ua_id", "value": "Yp8GoVBn7Tj0PWR7AAAAADaCnYPYpvql9WPFzFkbL_A=", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1788658652.910527, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "wxuin", "value": "53845090997782", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1788405090.13786, "httpOnly": false, "secure": false, "sameSite": "Lax"}, {"name": "uuid", "value": "8093d628ddf047d5a340ec437ca7fd08", "domain": "mp.weixin.qq.com", "path": "/", "expires": -1, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "rand_info", "value": "CAESIEjp7aANF1MLCho4lIEi6Ju8oYLEexqCUxELmbqpfG3Y", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754444283, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "slave_bizuin", "value": "3932822909", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754444283, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "data_bizuin", "value": "3932822909", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754444283, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "bizuin", "value": "3932822909", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754444283, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "data_ticket", "value": "TXXMQ0chAAhF8D9olYk9AhfbDAjnS2jHUUO8GHs0e45nt6P/zjA38GlY5q12SD0E", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754444283, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "slave_sid", "value": "OHl1NHRHZjd2N3poQTFDQ2d5WVBBWnJ5bm9obkZNVjZGank5S2dxSUMwblRqeFUwalBSNEpTbHlrMjZ6cUtxT2dRWk9NdElDcGRmUko1a1F4MENlUWFja0FvVjZDZnY0RXVkT0NiNUJnOVlITGQzOHRGYTBLUG10UEZzSzJhdWozakN1SUo5VXZRNmdkUGlz", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754444283, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "slave_user", "value": "gh_5624b93dc545", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754444283, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "xid", "value": "fedc4a4eeb3b124f5812d5b88b1265f4", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1788405116.292443, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "mm_lang", "value": "zh_CN", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1788405116.292574, "httpOnly": false, "secure": true, "sameSite": "Lax"}, {"name": "_clck", "value": "3932822909|1|fy4|0", "domain": ".qq.com", "path": "/", "expires": 1785634062, "httpOnly": false, "secure": false, "sameSite": "Lax"}, {"name": "_clsk", "value": "wa3u8e|1754098659834|3|1|mp.weixin.qq.com/weheat-agent/payload/record", "domain": ".qq.com", "path": "/", "expires": 1754185059, "httpOnly": false, "secure": false, "sameSite": "Lax"}], "origins": [{"origin": "https://mp.weixin.qq.com", "localStorage": [{"name": "USER_AGENT_BROWSER_MAJOR_VERSION", "value": "127"}, {"name": "loginMode", "value": "1"}, {"name": "adTransitionText", "value": "noAdText"}, {"name": "hasClick", "value": "[0]"}]}]}
\ No newline at end of file
diff --git a/network/android/sening_vdp/sening_vdp2 b/network/android/sening_vdp/sening_vdp2
deleted file mode 100644
index cc6ed95..0000000
--- a/network/android/sening_vdp/sening_vdp2
+++ /dev/null
Binary files differ
diff --git a/network/sening_vdp2.exe b/network/sening_vdp2.exe
deleted file mode 100644
index 0993b07..0000000
--- a/network/sening_vdp2.exe
+++ /dev/null
Binary files differ
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/__init__.py b/scripts/__init__.py
similarity index 100%
copy from auto_media_publisher/modifiers/video_matting/matting_model/__init__.py
copy to scripts/__init__.py
diff --git a/scripts/girls.json b/scripts/girls.json
new file mode 100644
index 0000000..a178b44
--- /dev/null
+++ b/scripts/girls.json
@@ -0,0 +1,99 @@
+[
+ "Beautiful East Asian woman in traditional hanfu, ornate hairpin, standing in a classical Chinese garden, ultra-realistic 8k photography",
+ "Elegant oriental woman in a red silk cheongsam, pearl hair ornaments, standing beside a lacquered wooden screen, cinematic lighting",
+ "Graceful woman in Tang dynasty attire, golden phoenix hair crown, posing in a bamboo grove, ultra-detailed realistic portrait",
+ "Beautiful young woman in flowing white hanfu, silver hairpin, standing on a stone bridge over a koi pond, soft daylight",
+ "Oriental beauty in embroidered qipao, jade earrings, seated at a tea table in a serene courtyard, high resolution",
+ "East Asian woman in Ming dynasty hanfu, holding a traditional paper umbrella, walking under cherry blossom trees",
+ "Elegant lady in blue silk hanfu, intricate floral headpiece, standing in a misty mountain landscape",
+ "Beautiful woman in pastel hanfu, butterfly hair ornaments, seated on a wooden veranda overlooking a lotus pond",
+ "Graceful oriental lady in traditional attire, golden tassel earrings, reading a book in an antique library",
+ "East Asian beauty in a red embroidered cheongsam, leaning against an ancient temple pillar",
+ "Beautiful woman in silk hanfu, white jade hairpin, holding a guqin in a candle-lit room",
+ "Elegant lady in purple hanfu, golden hair ornament, sitting in a moon gate courtyard",
+ "Graceful oriental beauty in green silk hanfu, standing beside a flowing river",
+ "Beautiful East Asian woman in a qipao with floral embroidery, sipping tea in a modern high-rise tea lounge",
+ "Young woman in Tang dynasty costume, delicate pearl headdress, standing in a palace corridor",
+ "Elegant lady in flowing hanfu, cherry blossom petals falling around her",
+ "Beautiful woman in modern city street wearing a retro qipao, neon lights in background",
+ "Graceful oriental lady in a white silk robe, walking barefoot on a tatami mat floor",
+ "Elegant woman in pink hanfu, golden phoenix crown, holding a lantern in a night market",
+ "Beautiful East Asian woman in hanfu, jade necklace, sitting beside a bonsai tree",
+ "Elegant lady in qipao, embroidered dragon design, walking through a rainy alley with an oil-paper umbrella",
+ "Graceful woman in pale blue hanfu, silver hairpin, gazing out from a tea house balcony",
+ "Beautiful lady in Tang dynasty gown, golden tassels in hair, standing by a lotus pond",
+ "Elegant woman in embroidered qipao, posing in front of a traditional wooden lattice window",
+ "Graceful lady in hanfu, holding a folding fan with painted plum blossoms",
+ "Beautiful woman in hanfu, hair adorned with red tassels, sitting on stone steps",
+ "Elegant lady in a high-collared cheongsam, pearl earrings, modern apartment interior",
+ "Graceful woman in flowing white hanfu, walking in a snowy courtyard",
+ "Elegant oriental lady in a light green hanfu, silk belt flowing in the wind",
+ "Beautiful woman in a Tang-style dress, standing in a candlelit banquet hall",
+ "Elegant lady in embroidered hanfu, surrounded by hanging paper lanterns",
+ "Graceful woman in a white silk qipao, leaning on a red lacquer table",
+ "Beautiful oriental lady in hanfu, standing at the entrance of a temple",
+ "Elegant woman in red cheongsam, holding a vintage Chinese folding fan",
+ "Graceful lady in traditional attire, gazing into a round moon gate",
+ "Beautiful woman in a hanfu with long flowing sleeves, spinning in a courtyard",
+ "Elegant lady in a pale yellow hanfu, pearl tassels swaying in her hair",
+ "Graceful woman in a qipao, holding a tea cup with delicate hands",
+ "Beautiful East Asian lady in white hanfu, standing in a bamboo forest",
+ "Elegant woman in Tang dynasty gown, standing on a stone bridge at sunset",
+ "Graceful lady in hanfu, sitting beside a guzheng under warm lamplight",
+ "Beautiful woman in embroidered silk robe, holding a scroll in a study",
+ "Elegant lady in red hanfu, golden hair crown, standing in a palace garden",
+ "Graceful woman in pastel hanfu, oil-paper umbrella resting on her shoulder",
+ "Beautiful lady in a blue silk qipao, standing near a koi pond",
+ "Elegant oriental woman in white hanfu, leaning against a cherry tree",
+ "Graceful lady in a dark red cheongsam, modern skyscrapers in background",
+ "Beautiful woman in a Tang-style dress, walking on palace marble steps",
+ "Elegant lady in flowing green hanfu, holding a pipa in a courtyard",
+ "Graceful woman in a floral qipao, standing by a red lantern-lit street",
+ "Beautiful East Asian lady in hanfu, gazing out a wooden lattice window",
+ "Elegant woman in white silk robe, holding incense in a temple",
+ "Graceful lady in a blue hanfu, walking on a misty riverside path",
+ "Beautiful woman in a Tang gown, holding a jade fan, standing in a pavilion",
+ "Elegant lady in a pale pink hanfu, holding a peach blossom branch",
+ "Graceful woman in qipao, standing at a balcony with a night city view",
+ "Beautiful lady in white hanfu, sitting on a rock by a flowing stream",
+ "Elegant oriental woman in silk robe, drinking tea under soft daylight",
+ "Graceful lady in red hanfu, leaning against a carved wooden railing",
+ "Beautiful woman in a Tang-style robe, standing under paper lanterns",
+ "Elegant lady in a light green qipao, walking in a flower garden",
+ "Graceful woman in embroidered hanfu, holding a brush and calligraphy paper",
+ "Beautiful lady in white hanfu, gazing into the distance on a hilltop",
+ "Elegant woman in Tang gown, sitting on a carved wooden chair",
+ "Graceful lady in a dark silk qipao, holding a crystal wine glass",
+ "Beautiful woman in a hanfu, resting by a window with lace curtains",
+ "Elegant lady in pink qipao, walking past traditional wooden houses",
+ "Graceful woman in blue hanfu, kneeling beside a stone lantern",
+ "Beautiful lady in Tang-style attire, standing under a blooming tree",
+ "Elegant oriental woman in white hanfu, surrounded by peach blossoms",
+ "Graceful lady in a qipao, sipping tea in a quiet courtyard",
+ "Beautiful woman in hanfu, holding a red silk ribbon in the wind",
+ "Elegant lady in Tang dress, standing on a wooden pier at sunset",
+ "Graceful woman in embroidered qipao, walking past neon-lit shop signs",
+ "Beautiful lady in hanfu, playing a guqin under candlelight",
+ "Elegant woman in white silk qipao, standing by a riverbank",
+ "Graceful lady in Tang-style gown, walking barefoot in a palace garden",
+ "Beautiful oriental woman in hanfu, holding a glowing lantern at night",
+ "Elegant lady in pastel hanfu, sitting on a swing in a courtyard",
+ "Graceful woman in a cheongsam, holding a traditional hand mirror",
+ "Beautiful lady in white hanfu, walking along a riverside bamboo path",
+ "Elegant woman in Tang gown, holding a fan with a painted crane",
+ "Graceful lady in red silk qipao, standing in a vintage tea house",
+ "Beautiful woman in hanfu, leaning on a carved stone balustrade",
+ "Elegant lady in light blue hanfu, oil-paper umbrella in hand",
+ "Graceful woman in qipao, sitting at a table with calligraphy tools",
+ "Beautiful lady in Tang gown, playing a pipa under moonlight",
+ "Elegant woman in hanfu, walking across a garden stone path",
+ "Graceful lady in silk qipao, gazing out from a modern glass balcony",
+ "Beautiful oriental woman in hanfu, surrounded by mist in the mountains",
+ "Elegant lady in red Tang gown, posing by a lacquered screen",
+ "Graceful woman in embroidered qipao, walking past a bamboo fence",
+ "Beautiful lady in a hanfu, holding a long silk ribbon in a dance pose",
+ "Elegant oriental woman in white hanfu, resting in a pavilion",
+ "Graceful lady in pale pink qipao, holding a single white flower",
+ "Beautiful woman in Tang-style gown, walking down palace steps",
+ "Elegant lady in hanfu, sitting at a tea table under hanging lanterns"
+]
\ No newline at end of file
diff --git a/scripts/office.json b/scripts/office.json
new file mode 100644
index 0000000..c2ae165
--- /dev/null
+++ b/scripts/office.json
@@ -0,0 +1,279 @@
+[
+ "A spacious modern corporate reception area with a large illuminated company sign, warm lighting, minimalist interior design, natural wood and marble finishes",
+ "An open, luxurious office lobby with glass walls and a sleek reception desk, large indoor plants, and soft daylight streaming in",
+ "A bright corporate entrance hall, featuring a bold wall logo, polished marble floor, and modern lighting fixtures",
+ "A large open-plan office with rows of desks, ergonomic chairs, modern computers, panoramic city view through floor-to-ceiling windows",
+ "A spacious tech company open workspace with glass partitions, indoor greenery, clean white desks, and a prominent wall logo",
+ "A minimalist open office interior, bright natural lighting, employees working on laptops, with a visible company sign in the background",
+ "A modern boardroom with a long wooden meeting table, leather executive chairs, a large screen displaying the company logo, surrounded by glass walls",
+ "A high-tech conference room, large video wall showing the company branding, neatly arranged chairs, and ceiling-mounted LED lights",
+ "A clean corporate meeting space with sunlight pouring in, digital projector displaying company presentation, and elegant d茅cor",
+ "A group of employees collaborating at a large table in a bright office, whiteboard with project notes, modern casual seating",
+ "Tech team having a brainstorming session in a creative workspace, sticky notes on the wall, glass panel with company branding behind them",
+ "Casual lounge area in a tech company office, employees chatting, coffee tables, decorative lighting",
+ "Close-up of a company reception desk nameplate, blurred background of an open office",
+ "Wall-mounted illuminated company logo in a corporate hallway",
+ "Office window view with reflections of the city skyline on glass walls",
+ "A bright open workspace with clusters of modern workstations, comfortable chairs, and greenery",
+ "A collaborative co-working style office, large communal desks, pendant lights, and modern furnishings",
+ "A luxurious executive office with a panoramic city view, leather chairs, and a large desk",
+ "A creative design studio office with colorful furniture, brainstorming boards, and open shelves",
+ "A glass-walled meeting pod inside a modern open-plan office, surrounded by plants and natural light",
+ "Spacious modern office lobby with marble floors and glass walls, bright natural light",
+ "Open-plan tech office with long desks, ergonomic chairs, and floor-to-ceiling windows",
+ "Reception area with a sleek desk, potted plants, and a backlit company sign",
+ "Corporate boardroom with a long wooden table, leather chairs, and a panoramic city view",
+ "Bright meeting room with glass walls, large monitor, and modern pendant lighting",
+ "Creative workspace with colorful furniture, open shelving, and brainstorming boards",
+ "Casual breakout area with comfortable sofas and coffee tables",
+ "Private office with a large desk, executive chair, and bookshelves",
+ "Large open workspace with clustered desks and indoor plants",
+ "High-tech conference room with video conferencing equipment and wall-mounted screens",
+ "Reception lobby with a minimalist design, wood accents, and hanging lights",
+ "Co-working space with shared desks, pendant lights, and communal seating",
+ "Hallway with framed artwork and glass partitions leading to office rooms",
+ "Modern lounge area with plush armchairs and low coffee tables",
+ "Open office with exposed brick walls and industrial-style lighting",
+ "Corporate cafeteria with long communal tables and a self-service coffee station",
+ "Bright workspace with natural wood desks and greenery",
+ "Minimalist reception area with a monochrome color palette",
+ "Casual collaborative zone with beanbags and whiteboards",
+ "Elegant meeting room with floor-length curtains and a large chandelier",
+ "Tech startup workspace with exposed ceilings and whiteboard walls",
+ "Open office floor with large computer monitors and ergonomic setups",
+ "Quiet reading corner with lounge chairs and bookshelves",
+ "Executive private office with glass walls and a large city skyline view",
+ "Training room with rows of desks and a projector screen",
+ "Open seating area with round tables and small plants",
+ "Modern conference table surrounded by sleek black chairs",
+ "Office library with tall bookshelves and soft armchairs",
+ "Design team workspace with corkboards and sticky notes",
+ "Reception area with marble reception desk and fresh flowers",
+ "Glass-enclosed private meeting pod in an open office",
+ "Casual staff kitchen with bar stools and coffee machines",
+ "Large collaboration table with laptops and notebooks",
+ "Industrial-style workspace with exposed pipes and concrete floors",
+ "Colorful creative hub with neon signs and mood boards",
+ "Quiet work booth with acoustic panels for privacy",
+ "Panoramic open office with city skyline visible through large windows",
+ "Cozy break room with sofa seating and snack shelves",
+ "Shared workspace with long communal workbenches",
+ "Glass-walled huddle room with small round table",
+ "Conference room with dual monitors and teleconferencing gear",
+ "Office lounge with wall art and low lighting",
+ "Open cubicle layout with semi-transparent dividers",
+ "Minimalist standing desks with laptops and clean surfaces",
+ "Reception lobby with floor plants and pendant lighting",
+ "Shared hot-desk area with open shelves",
+ "Large presentation room with tiered seating",
+ "Bright brainstorming space with writable glass walls",
+ "Staff dining area with modern tables and chairs",
+ "Executive boardroom with leather high-back chairs",
+ "Tech company workspace with wall-mounted whiteboards",
+ "Reception desk backed by a textured feature wall",
+ "Creative meeting space with low sofas and round tables",
+ "Project room with large pin-up boards and sticky notes",
+ "Pantry area with coffee machines and microwaves",
+ "Casual team meeting in a sunlit lounge",
+ "Cubicle workstations with desktop computers",
+ "Office with wide corridors and glass meeting rooms",
+ "Large table in the center of an open workspace",
+ "Quiet lounge with potted plants and soft armchairs",
+ "Collaboration corner with movable whiteboards",
+ "Minimal office layout with clean desks and storage cabinets",
+ "Reception area with feature lighting and natural wood",
+ "Breakout zone with colorful stools and round tables",
+ "Workshop-style office with large drafting tables",
+ "Boardroom with long oval table and built-in monitor",
+ "Open floor with casual caf茅-style seating",
+ "Training center with projector and student desks",
+ "Casual reading nook with beanbags and bookshelves",
+ "Panoramic executive office with wall-to-wall windows",
+ "Informal collaboration space with high tables and stools",
+ "Reception hall with wide marble staircase",
+ "Conference room with modern art on the walls",
+ "Bright office lounge with colorful armchairs",
+ "Minimal workspace with white desks and sleek chairs",
+ "Quiet private office with frosted glass partitions",
+ "Corporate lobby with tall ceilings and skylights",
+ "Collaborative tech lab with monitors and cables",
+ "Pantry corner with bar counter and stools",
+ "Creative office corner with easels and sketch boards",
+ "Open-plan work area with soft carpeting",
+ "Small huddle room with round white table",
+ "Reception desk in front of a plant wall",
+ "Bright open space with low partition desks",
+ "Boardroom with ceiling-mounted projector",
+ "Office dining lounge with small round tables",
+ "Hallway with lockers and office artwork",
+ "Large atrium-style office lobby",
+ "Cubicle section with task lighting",
+ "Open area with large group worktables",
+ "Modern office with indoor tree planter",
+ "Quiet office booth with built-in seating",
+ "Casual caf茅 corner in the workplace",
+ "Reception lounge with a fireplace",
+ "Tech workspace with VR headsets",
+ "Private interview room with two chairs",
+ "Bright hallway leading to open workspace",
+ "Training hall with large display screens",
+ "Open-plan office with high ceilings and exposed beams",
+ "Reception with hanging pendant lights and terrazzo floors",
+ "Executive lounge with leather sofas and coffee tables",
+ "modern open-plan tech office, wide space, glass walls, natural sunlight streaming in, ergonomic desks, no people",
+ "modern open-plan tech office, wide space, glass walls, natural sunlight streaming in, ergonomic desks, busy employees collaborating",
+ "bright coworking space in a tech startup, shared desks, wide hallways, plants, no people",
+ "bright coworking space in a tech startup, shared desks, wide hallways, plants, young team brainstorming",
+ "minimalist technology company office, clean lines, large windows, empty workspace",
+ "minimalist technology company office, clean lines, large windows, group meeting at a long table",
+ "wide panoramic view of a futuristic tech office, sleek furniture, warm lighting, no people",
+ "wide panoramic view of a futuristic tech office, sleek furniture, warm lighting, engineers working at laptops",
+ "spacious startup office, exposed brick walls, industrial lights, modern desks, no people",
+ "spacious startup office, exposed brick walls, industrial lights, modern desks, software developers at work",
+ "modern office lobby in a tech company, glass doors, digital screens, no people",
+ "modern office lobby in a tech company, glass doors, digital screens, reception staff greeting visitors",
+ "high-tech innovation lab, open workstations, LED lighting, wide angles, no people",
+ "high-tech innovation lab, open workstations, LED lighting, team testing prototypes",
+ "large open-plan office with adjustable standing desks, bright daylight, no people",
+ "large open-plan office with adjustable standing desks, bright daylight, employees coding",
+ "corporate tech headquarters, wide hallway, glass partitions, empty meeting rooms",
+ "corporate tech headquarters, wide hallway, glass partitions, team members chatting",
+ "modern coworking space with vibrant colors, wide desks, no people",
+ "modern coworking space with vibrant colors, wide desks, young team collaborating on whiteboards",
+ "open-plan technology company office, natural wood desks, large green plants, no people",
+ "open-plan technology company office, natural wood desks, large green plants, designers working on laptops",
+ "startup workspace with long communal tables, wide airy feel, no people",
+ "startup workspace with long communal tables, wide airy feel, small team in discussion",
+ "bright and airy product design studio, high ceilings, modern furniture, no people",
+ "bright and airy product design studio, high ceilings, modern furniture, designers reviewing prototypes",
+ "modern UX design office, spacious, glass walls, minimal decor, no people",
+ "modern UX design office, spacious, glass walls, minimal decor, team sketching wireframes",
+ "spacious game development studio, multiple monitors on desks, LED ambient lighting, no people",
+ "spacious game development studio, multiple monitors on desks, LED ambient lighting, developers playtesting",
+ "high-ceiling innovation space, collaborative seating, no people",
+ "high-ceiling innovation space, collaborative seating, creative team brainstorming",
+ "large-scale open-plan tech office, quiet empty morning scene",
+ "large-scale open-plan tech office, busy afternoon with employees at work",
+ "modern engineering office, long desks, white walls, no people",
+ "modern engineering office, long desks, white walls, engineers soldering electronics",
+ "creative startup space, beanbags, whiteboards, colorful decor, no people",
+ "creative startup space, beanbags, whiteboards, colorful decor, brainstorming session",
+ "modern fintech company office, minimalist design, empty workspace",
+ "modern fintech company office, minimalist design, financial analysts at work",
+ "AI research lab office, futuristic furniture, no people",
+ "AI research lab office, futuristic furniture, scientists discussing algorithms",
+ "wide corporate innovation hub, bright daylight, no people",
+ "wide corporate innovation hub, bright daylight, networking event with tech founders",
+ "large open-plan office with meeting pods, no people",
+ "large open-plan office with meeting pods, startup founders having a meeting",
+ "open-concept robotics company office, tool benches, wide space, no people",
+ "open-concept robotics company office, tool benches, engineers working on robots",
+ "wide photography of cloud computing company office, bright and airy, no people",
+ "wide photography of cloud computing company office, bright and airy, employees at their desks",
+ "cozy corner in a tech office, couch seating, natural light, no people",
+ "cozy corner in a tech office, couch seating, natural light, small team chatting casually",
+ "wide meeting room in a technology firm, modern chairs, no people",
+ "wide meeting room in a technology firm, modern chairs, executives discussing strategy",
+ "creative media tech company workspace, open desks, natural wood, no people",
+ "creative media tech company workspace, open desks, natural wood, designers editing videos",
+ "startup incubator open floor, colorful chairs, no people",
+ "startup incubator open floor, colorful chairs, entrepreneurs networking",
+ "modern VR/AR development office, large space, no people",
+ "modern VR/AR development office, large space, developers testing VR headsets",
+ "wide hallway in a tech firm, glass conference rooms, no people",
+ "wide hallway in a tech firm, glass conference rooms, teams in meetings",
+ "innovative workspace with writable glass walls, no people",
+ "innovative workspace with writable glass walls, team sketching app ideas",
+ "modern coworking lounge area, bright colors, wide sofa seating, no people",
+ "modern coworking lounge area, bright colors, wide sofa seating, group discussion",
+ "spacious hackathon event space, open desks, no people",
+ "spacious hackathon event space, open desks, participants coding together",
+ "corporate training room in a tech firm, wide desks, projectors, no people",
+ "corporate training room in a tech firm, wide desks, projectors, employees attending a session",
+ "open space in a game design studio, colorful lighting, no people",
+ "open space in a game design studio, colorful lighting, game designers discussing concepts",
+ "wide cafeteria in a tech headquarters, modern design, no people",
+ "wide cafeteria in a tech headquarters, modern design, employees having lunch",
+ "startup collaboration hub, high tables, wide view, no people",
+ "startup collaboration hub, high tables, wide view, young team having coffee",
+ "modern AI startup office, sleek decor, no people",
+ "modern AI startup office, sleek decor, machine learning engineers collaborating",
+ "spacious tech workshop, 3D printers, open benches, no people",
+ "spacious tech workshop, 3D printers, open benches, makers working",
+ "innovative design lab, bright lighting, no people",
+ "innovative design lab, bright lighting, design team working on prototypes",
+ "wide open-plan office with ceiling skylights, no people",
+ "wide open-plan office with ceiling skylights, developers collaborating",
+ "startup brainstorming space, writable walls, high stools, no people",
+ "startup brainstorming space, writable walls, high stools, creative team sketching ideas",
+ "corporate innovation lounge, sleek sofas, no people",
+ "corporate innovation lounge, sleek sofas, small team meeting casually",
+ "AI company meeting space, futuristic furniture, no people",
+ "AI company meeting space, futuristic furniture, executives discussing strategy",
+ "open-plan space in a green tech startup, plants everywhere, no people",
+ "open-plan space in a green tech startup, plants everywhere, employees coding",
+ "modern minimalist workspace, clean desks, no people",
+ "modern minimalist workspace, clean desks, focused programmers at work",
+ "large hallway in a software company, bright daylight, no people",
+ "large hallway in a software company, bright daylight, employees walking",
+ "VR gaming startup office, neon lighting, wide desks, no people",
+ "VR gaming startup office, neon lighting, wide desks, developers testing VR games",
+ "open collaboration area in a creative tech firm, colorful chairs, no people",
+ "open collaboration area in a creative tech firm, colorful chairs, group discussion",
+ "corporate innovation center atrium, skylight, no people",
+ "corporate innovation center atrium, skylight, employees networking",
+ "spacious developer workspace, dual monitors on every desk, no people",
+ "spacious developer workspace, dual monitors on every desk, programmers typing",
+ "clean open-concept office with white furniture, no people",
+ "clean open-concept office with white furniture, software engineers collaborating",
+ "wide team room in a software company, large table, no people",
+ "wide team room in a software company, large table, sprint planning session",
+ "innovative product lab with prototypes on tables, no people",
+ "innovative product lab with prototypes on tables, engineers testing gadgets",
+ "corporate lounge with coffee tables, wide space, no people",
+ "corporate lounge with coffee tables, wide space, employees chatting",
+ "open-plan design office with drafting tables, no people",
+ "open-plan design office with drafting tables, industrial designers at work",
+ "large reception area in a technology company, wide view, no people",
+ "large reception area in a technology company, wide view, receptionist greeting guests",
+ "creative hack space in a startup hub, tools on walls, no people",
+ "creative hack space in a startup hub, tools on walls, makers collaborating",
+ "modern meeting pod in an open office, no people",
+ "modern meeting pod in an open office, two employees having a conversation",
+ "innovation lab with robotics on tables, no people",
+ "innovation lab with robotics on tables, engineers assembling machines",
+ "spacious coding floor, rows of desks, no people",
+ "spacious coding floor, rows of desks, developers coding together",
+ "corporate strategy room, wall screens, no people",
+ "corporate strategy room, wall screens, executives reviewing analytics",
+ "open seating lounge in a tech incubator, wide view, no people",
+ "open seating lounge in a tech incubator, wide view, startup founders talking",
+ "spacious cafeteria with modern lighting, no people",
+ "spacious cafeteria with modern lighting, employees having casual chats",
+ "wide break area with coffee machines, no people",
+ "wide break area with coffee machines, employees enjoying coffee",
+ "large open floor for software engineers, no people",
+ "large open floor for software engineers, engineers pair programming",
+ "tech company boardroom, modern minimalist design, no people",
+ "tech company boardroom, modern minimalist design, leadership meeting",
+ "AI research open workspace, futuristic screens, no people",
+ "AI research open workspace, futuristic screens, scientists coding",
+ "corporate innovation center, clean lines, spacious, no people",
+ "corporate innovation center, clean lines, spacious, networking session",
+ "product development lab, bright daylight, no people",
+ "product development lab, bright daylight, engineers testing devices",
+ "spacious coworking cafe area, no people",
+ "spacious coworking cafe area, employees chatting casually",
+ "startup hot desk zone, clean and bright, no people",
+ "startup hot desk zone, clean and bright, workers with laptops",
+ "modern hallway with glass meeting pods, no people",
+ "modern hallway with glass meeting pods, people having discussions inside",
+ "spacious innovation floor with colorful seating, no people",
+ "spacious innovation floor with colorful seating, team collaboration session",
+ "large bright conference space, no people",
+ "large bright conference space, panel discussion event",
+ "modern office balcony with city view, no people",
+ "modern office balcony with city view, employees relaxing",
+ "tech company innovation hub with product displays, no people",
+ "tech company innovation hub with product displays, visitors exploring"
+]
\ No newline at end of file
diff --git a/scripts/work_gen_images.py b/scripts/work_gen_images.py
new file mode 100644
index 0000000..734c1e3
--- /dev/null
+++ b/scripts/work_gen_images.py
@@ -0,0 +1,46 @@
+
+
+from auto_media_publisher.utils.utils_text2image_pollinations import Text2ImageGenerator
+from auto_media_publisher.utils.utils_file import load_json
+import os
+import random
+import datetime
+import sys
+import time
+
+def main():
+
+ count = 100
+
+ if len(sys.argv) >1:
+ try:
+ count = int(sys.argv[1])
+ if count >100 or count < 1:
+ count =100
+ except:
+ pass
+
+ generator = Text2ImageGenerator()
+
+ # 寰呯敓鎴愮殑涓枃鎻愮ず璇嶏紙蹇呭～锛�
+ prompts = load_json("./scripts/office.json")
+
+ output_dir = "./scripts/output"
+ os.makedirs(output_dir,exist_ok=True)
+ # 鐢熸垚骞朵繚瀛樺浘鍍�
+ try:
+ for i in range(count):
+ prompt = random.choice(prompts)
+ now = datetime.datetime.now()
+ filename = os.path.join(output_dir,now.strftime("%Y%m%d%H%M%S") + ".png")
+
+ saved_path = generator.generate_image_without_trans(prompt, filename=filename)
+ print(f"[鈭歖 鍥剧墖鎴愬姛淇濆瓨鑷筹細{saved_path}")
+ time.sleep(random.randint(60,80))
+
+ except Exception as e:
+ print(f"[脳] 鐢熸垚澶辫触锛歿e}")
+
+if __name__ == "__main__":
+ main()
+
diff --git a/test/test_face_swap.py b/test/test_face_swap.py
index 675a059..7825cef 100644
--- a/test/test_face_swap.py
+++ b/test/test_face_swap.py
@@ -12,7 +12,7 @@
start_time = time.time()
- processor.swap("./temp/buhui.mp4", "./temp/buhui-lilizhen-with-new.mp4","./temp/lilizhen.jpg")
+ processor.swap("./temp/dancing.mp4", "./temp/dancing-lilizhen-with-new.mp4","./temp/lilizhen.jpg")
print(time.time()-start_time)
diff --git a/test/test_matanyone.py b/test/test_matanyone.py
new file mode 100644
index 0000000..d41f623
--- /dev/null
+++ b/test/test_matanyone.py
@@ -0,0 +1,25 @@
+
+
+from auto_media_publisher.modifiers.video_matting.mat_anyone import MatAnyoneConverter
+import time
+if __name__ == '__main__':
+
+ converter = MatAnyoneConverter()
+
+
+ input = './temp/dancing.mp4'
+ output = './temp'
+
+ bg_image = "./temp/background.png"
+ bg_video = "./temp/duoren.mp4"
+
+ time_start = time.time()
+
+
+ #converter.video_matting(input,output,None,None)
+ converter.video_matting(input,output,bg_video,None)
+ print(f"video matting finished .. {time.time()-time_start}")
+ # conda activate mediapublish
+ # f:
+ # cd automedia
+ #python -m test.test_video_matting --input ./temp/buhui.mp4 --bg_video ./temp/ocean.mp4 --output_video ./temp/buhui_ocean-36.mp4
diff --git a/test/test_text2image.py b/test/test_text2image.py
index a2696af..31dae76 100644
--- a/test/test_text2image.py
+++ b/test/test_text2image.py
@@ -12,7 +12,7 @@
# 鐢熸垚骞朵繚瀛樺浘鍍�
try:
- saved_path = generator.generate_image(prompt, filename=output_file)
+ saved_path = generator.generate_image_with_trans(prompt, filename=output_file)
print(f"[鈭歖 鍥剧墖鎴愬姛淇濆瓨鑷筹細{saved_path}")
except Exception as e:
print(f"[脳] 鐢熸垚澶辫触锛歿e}")
diff --git a/test/test_video_matting.py b/test/test_video_matting.py
index 7d9e8f4..6bb5700 100644
--- a/test/test_video_matting.py
+++ b/test/test_video_matting.py
@@ -1,6 +1,6 @@
-from auto_media_publisher.modifiers.video_matting.video_matting_converter import VideoMattingConverter
+from auto_media_publisher.modifiers.video_matting.robust_video_matting_converter import RobustVideoMattingConverter
import time
if __name__ == '__main__':
import argparse
@@ -8,7 +8,7 @@
parser = argparse.ArgumentParser()
parser.add_argument('--input', type=str, required=True)
parser.add_argument('--bg_video', type=str)
- parser.add_argument('--bg_jpg', type=str)
+ parser.add_argument('--bg_image', type=str)
parser.add_argument('--output_video', type=str)
parser.add_argument('--output_alpha', type=str)
parser.add_argument('--output_fore', type=str)
@@ -17,7 +17,7 @@
time_start = time.time()
- converter = VideoMattingConverter()
+ converter = RobustVideoMattingConverter()
print(f"model loaded .. {time.time()-time_start}")
time_start = time.time()
@@ -25,7 +25,7 @@
converter.video_matting(
input=args.input,
bg_video=args.bg_video,
- bg_jpg=args.bg_jpg,
+ bg_image=args.bg_image,
output_video=args.output_video,
output_alpha=args.output_alpha,
output_fore=args.output_fore,
--
Gitblit v1.9.1