comi-logic/Wan2GP
8
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merged WanGP4 features

Browse Source
This commit is contained in:
DeepBeepMeep
2025-04-15 13:49:07 +02:00
parent 0ea781ea12 c8a538b42a
commit d4a498edb3
81 changed files with 14682 additions and 3622 deletions
Download Patch File Download Diff File Expand all files Collapse all files

92
README.md
View File

@@ -10,15 +10,19 @@
<b>Wan2.1 GP by DeepBeepMeep based on Wan2.1's Alibaba: Open and Advanced Large-Scale Video Generative Models for the GPU Poor</b>
</p>
In this repository, we present **Wan2.1**, a comprehensive and open suite of video foundation models that pushes the boundaries of video generation. **Wan2.1** offers these key features:
- 👍 **SOTA Performance**: **Wan2.1** consistently outperforms existing open-source models and state-of-the-art commercial solutions across multiple benchmarks.
- 👍 **Supports Consumer-grade GPUs**: The T2V-1.3B model requires only 8.19 GB VRAM, making it compatible with almost all consumer-grade GPUs. It can generate a 5-second 480P video on an RTX 4090 in about 4 minutes (without optimization techniques like quantization). Its performance is even comparable to some closed-source models.
- 👍 **Multiple Tasks**: **Wan2.1** excels in Text-to-Video, Image-to-Video, Video Editing, Text-to-Image, and Video-to-Audio, advancing the field of video generation.
- 👍 **Visual Text Generation**: **Wan2.1** is the first video model capable of generating both Chinese and English text, featuring robust text generation that enhances its practical applications.
- 👍 **Powerful Video VAE**: **Wan-VAE** delivers exceptional efficiency and performance, encoding and decoding 1080P videos of any length while preserving temporal information, making it an ideal foundation for video and image generation.
## 🔥 Latest News!!
* April 13 2025: 👋 Wan 2.1GP v4.0: lots of goodies for you !
- A new UI, tabs were replaced by a Dropdown box to easily switch models
- A new queuing system that lets you stack in a queue as many text2video, imag2video tasks, ... as you want. Each task can rely on complete different generation parameters (different number of frames, steps, loras, ...). Many thanks to *Tophness** for being a big contributor on this new feature
- Temporal upsampling (Rife) and spatial upsampling (Lanczos) for a smoother video (32 fps or 64 fps) and to enlarge your video by x2 or x4. Check these new advanced options.
- Wan Vace Control Net support : with Vace you can inject in the scene people or objects, animate a person, perform inpainting or outpainting, continue a video, ... I have provided an introduction guide below.
- Integrated *Matanyone* tool directly inside WanGP so that you can create easily inpainting masks used in Vace
- Sliding Window generation for Vace, create windows that can last dozen of seconds
- New optimisations for old generation GPUs: Generate 5s (81 frames, 15 steps) of Vace 1.3B with only 5GB and in only 6 minutes on a RTX 2080Ti and 5s of t2v 14B in less than 10 minutes.
* Mar 27 2025: 👋 Added support for the new Wan Fun InP models (image2video). The 14B Fun InP has probably better end image support but unfortunately existing loras do not work so well with it. The great novelty is the Fun InP image2 1.3B model : Image 2 Video is now accessible to even lower hardware configuration. It is not as good as the 14B models but very impressive for its size. You can choose any of those models in the Configuration tab. Many thanks to the VideoX-Fun team (https://github.com/aigc-apps/VideoX-Fun)
* Mar 26 2025: 👋 Good news ! Official support for RTX 50xx please check the installation instructions below.
* Mar 24 2025: 👋 Wan2.1GP v3.2:
@@ -175,35 +179,35 @@ pip install -e .
To run the text to video generator (in Low VRAM mode):
```bash
python gradio_server.py
python wgp.py.py
#or
python gradio_server.py --t2v #launch the default text 2 video model
python wgp.py.py --t2v #launch the default text 2 video model
#or
python gradio_server.py --t2v-14B #for the 14B model
python wgp.py.py --t2v-14B #for the 14B model
#or
python gradio_server.py --t2v-1-3B #for the 1.3B model
python wgp.py.py --t2v-1-3B #for the 1.3B model
```
To run the image to video generator (in Low VRAM mode):
```bash
python gradio_server.py --i2v
python wgp.py.py --i2v
```
To run the 1.3B Fun InP image to video generator (in Low VRAM mode):
```bash
python gradio_server.py --i2v-1-3B
python wgp.py.py --i2v-1-3B
```
To be able to input multiple images with the image to video generator:
```bash
python gradio_server.py --i2v --multiple-images
python wgp.py.py --i2v --multiple-images
```
Within the application you can configure which video generator will be launched without specifying a command line switch.
To run the application while loading entirely the diffusion model in VRAM (slightly faster but requires 24 GB of VRAM for a 8 bits quantized 14B model )
```bash
python gradio_server.py --profile 3
python wgp.py.py --profile 3
```
**Trouble shooting**:\
@@ -216,7 +220,7 @@ Therefore you may have no choice but to fallback to sdpa attention, to do so:
or
- Launch the application this way:
```bash
python gradio_server.py --attention sdpa
python wgp.py.py --attention sdpa
```
### Loras support
@@ -224,7 +228,7 @@ python gradio_server.py --attention sdpa
Every lora stored in the subfoler 'loras' for t2v and 'loras_i2v' will be automatically loaded. You will be then able to activate / desactive any of them when running the application by selecting them in the area below "Activated Loras" .
If you want to manage in differenta areas Loras for the 1.3B model and the 14B as they are not comptatible, just create the following subfolders:
If you want to manage in different areas Loras for the 1.3B model and the 14B as they are not compatible, just create the following subfolders:
- loras/1.3B
- loras/14B
@@ -250,7 +254,7 @@ Each preset, is a file with ".lset" extension stored in the loras directory and
Last but not least you can pre activate Loras corresponding and prefill a prompt (comments only or full prompt) by specifying a preset when launching the gradio server:
```bash
python gradio_server.py --lora-preset mylorapreset.lset # where 'mylorapreset.lset' is a preset stored in the 'loras' folder
python wgp.py.py --lora-preset mylorapreset.lset # where 'mylorapreset.lset' is a preset stored in the 'loras' folder
```
You will find prebuilt Loras on https://civitai.com/ or you will be able to build them with tools such as kohya or onetrainer.
@@ -271,6 +275,49 @@ In the video, a woman is presented. The woman is in a city and looks at her watc
You can define multiple lines of macros. If there is only one macro line, the app will generate a simple user interface to enter the macro variables when getting back to *Normal Mode* (advanced mode turned off)
### VACE ControlNet introduction
Vace is a ControlNet 1.3B text2video model that allows you to do Video to Video and Reference to Video (inject your own images into the output video). So with Vace you can inject in the scene people or objects of your choice, animate a person, perform inpainting or outpainting, continue a video, ...
First you need to select the Vace 1.3B model in the Drop Down box at the top. Please note that Vace works well for the moment only with videos up to 5s (81 frames).
Beside the usual Text Prompt, three new types of visual hints can be provided (and combined !):
- a Control Video: Based on your choice, you can decide to transfer the motion, the depth in a new Video. You can tell WanGP to use only the first n frames of Control Video and to extrapolate the rest. You can also do inpainting ). If the video contains area of grey color 127, they will be considered as masks and will be filled based on the Text prompt of the reference Images.
- reference Images: Use this to inject people or objects of your choice in the video. You can select multiple reference Images. The integration of the image is more efficient if the background is replaced by the full white color. You can do that with your preferred background remover or use the built in background remover by checking the box *Remove background*
- a Video Mask
This offers a stronger mechanism to tell Vace which parts should be kept (black) or replaced (white). You can do as well inpainting / outpainting, fill the missing part of a video more efficientlty with just the video hint. If a video mask is white, it will be generated so with black frames at the beginning and at the end and the rest white, you could generate the missing frames in between.
Examples:
- Inject people and / objects into a scene describe by a text prompt: Ref. Images + text Prompt
- Animate a character described in a text prompt: a Video of person moving + text Prompt
- Animate a character of your choice (motion transfer) : Ref Images + a Video of person moving + text Prompt
- Change the style of a scene (depth transfer): a Video that contains objects / person at differen depths + text Prompt
There are lots of possible combinations. Some of them require to prepare some materials (masks on top of video, full masks, etc...).
Vace provides on its github (https://github.com/ali-vilab/VACE/tree/main/vace/gradios) annotators / preprocessors Gradio tool that can help you build some of these materials depending on the task you want to achieve.
There is also a guide that describes the various combination of hints (https://github.com/ali-vilab/VACE/blob/main/UserGuide.md).Good luck !
It seems you will get better results if you turn on "Skip Layer Guidance" with its default configuration
### VACE Slidig Window
With this mode (that works for the moment only with Vace) you can merge mutiple Videos to form a very long video (up to 1 min). What is this very nice a about this feature is that the resulting video can be driven by the same control video. For instance the first 0-4s of the control video will be used to generate the first window then the next 4-8s of the control video will be used to generate the second window, and so on. So if your control video contains a person walking, your generate video could contain up to one minute of this person walking.
To turn on sliding window, you need to go in the Advanced Settings Tab *Sliding Window* and set the iteration number to a number greater than 1. This number corresponds to the default number of windows. You can still increase the number during the genreation by clicking the "One More Sample, Please !" button.
Each window duration will be set by the *Number of frames (16 = 1s)* form field. However the actual number of frames generated by each iteration will be less, because the *overlap frames* and *discard last frames*:
- *overlap frames* : the first frames ofa new window are filled with last frames of the previous window in order to ensure continuity between the two windows
- *discard last frames* : quite often the last frames of a window have a worse quality. You decide here how many ending frames of a new window should be dropped.
Number of Generated = [Number of iterations] * ([Number of frames] - [Overlap Frames] - [Discard Last Frames]) + [Overlap Frames]
Experimental: if your prompt is broken into multiple lines (each line separated by a carriage return), then each line of the prompt will be used for a new window. If there are more windows to generate than prompt lines, the last prompt line will be repeated.
### Command line parameters for Gradio Server
--i2v : launch the image to video generator\
--t2v : launch the text to video generator (default defined in the configuration)\
@@ -278,6 +325,7 @@ You can define multiple lines of macros. If there is only one macro line, the ap
--t2v-1-3B : launch the 1.3B model text to video generator\
--i2v-14B : launch the 14B model image to video generator\
--i2v-1-3B : launch the Fun InP 1.3B model image to video generator\
--vace : launch the Vace ControlNet 1.3B model image to video generator\
--quantize-transformer bool: (default True) : enable / disable on the fly transformer quantization\
--lora-dir path : Path of directory that contains Loras in diffusers / safetensor format\
--lora-preset preset : name of preset gile (without the extension) to preload
@@ -291,7 +339,7 @@ You can define multiple lines of macros. If there is only one macro line, the ap
--compile : turn on pytorch compilation\
--attention mode: force attention mode among, sdpa, flash, sage, sage2\
--profile no : default (4) : no of profile between 1 and 5\
--preload no : number in Megabytes to preload partially the diffusion model in VRAM , may offer slight speed gains especially on older hardware. Works only with profile 2 and 4.\
--preload no : number in Megabytes to preload partially the diffusion model in VRAM , may offer speed gains on older hardware, on recent hardware (RTX 30XX, RTX40XX and RTX50XX) speed gain is only 10% and not worth it. Works only with profile 2 and 4.\
--seed no : set default seed value\
--frames no : set the default number of frames to generate\
--steps no : set the default number of denoising steps\
@@ -299,10 +347,12 @@ You can define multiple lines of macros. If there is only one macro line, the ap
--slg : turn on skip layer guidance for improved quality\
--check-loras : filter loras that are incompatible (will take a few seconds while refreshing the lora list or while starting the app)\
--advanced : turn on the advanced mode while launching the app\
--i2v-settings : path to launch settings for i2v\
--t2v-settings : path to launch settings for t2v\
--listen : make server accessible on network\
--gpu device : run Wan on device for instance "cuda:1"
--gpu device : run Wan on device for instance "cuda:1"\
--settings: path a folder that contains the default settings for all the models\
--fp16: force to use fp16 versions of models instead of bf16 versions\
--perc-reserved-mem-max float_less_than_1 : max percentage of RAM to allocate to reserved RAM, allow faster transfers RAM<->VRAM. Value should remain below 0.5 to keep the OS stable\
--theme theme_name: load the UI with the specified Theme Name, so far only two are supported, "default" and "gradio". You may submit your own nice looking Gradio theme and I will add them
### Profiles (for power users only)
You can choose between 5 profiles, but two are really relevant here :

306
gradio/i2v_14B_singleGPU.py
View File

@@ -1,306 +0,0 @@
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import argparse
import gc
import os.path as osp
import os
import sys
import warnings
import gradio as gr
warnings.filterwarnings('ignore')
# Model
sys.path.insert(0, os.path.sep.join(osp.realpath(__file__).split(os.path.sep)[:-2]))
import wan
from wan.configs import MAX_AREA_CONFIGS, WAN_CONFIGS
from wan.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander
from wan.utils.utils import cache_video
# Global Var
prompt_expander = None
wan_i2v_480P = None
wan_i2v_720P = None
# Button Func
def load_i2v_model(value):
global wan_i2v_480P, wan_i2v_720P
from mmgp import offload
if value == '------':
print("No model loaded")
return '------'
if value == '720P':
if args.ckpt_dir_720p is None:
print("Please specify the checkpoint directory for 720P model")
return '------'
if wan_i2v_720P is not None:
pass
else:
del wan_i2v_480P
gc.collect()
wan_i2v_480P = None
print("load 14B-720P i2v model...", end='', flush=True)
cfg = WAN_CONFIGS['i2v-14B']
wan_i2v_720P = wan.WanI2V(
config=cfg,
checkpoint_dir=args.ckpt_dir_720p,
device_id=0,
rank=0,
t5_fsdp=False,
dit_fsdp=False,
use_usp=False,
i2v720p= True
)
print("done", flush=True)
pipe = {"transformer": wan_i2v_720P.model, "text_encoder" : wan_i2v_720P.text_encoder.model, "text_encoder_2": wan_i2v_720P.clip.model, "vae": wan_i2v_720P.vae.model } #
offload.profile(pipe, profile_no=4, budgets = {"transformer":100, "*":3000}, verboseLevel=2, compile="transformer", quantizeTransformer = False, pinnedMemory = False)
return '720P'
if value == '480P':
if args.ckpt_dir_480p is None:
print("Please specify the checkpoint directory for 480P model")
return '------'
if wan_i2v_480P is not None:
pass
else:
del wan_i2v_720P
gc.collect()
wan_i2v_720P = None
print("load 14B-480P i2v model...", end='', flush=True)
cfg = WAN_CONFIGS['i2v-14B']
wan_i2v_480P = wan.WanI2V(
config=cfg,
checkpoint_dir=args.ckpt_dir_480p,
device_id=0,
rank=0,
t5_fsdp=False,
dit_fsdp=False,
use_usp=False,
i2v720p= False
)
print("done", flush=True)
pipe = {"transformer": wan_i2v_480P.model, "text_encoder" : wan_i2v_480P.text_encoder.model, "text_encoder_2": wan_i2v_480P.clip.model, "vae": wan_i2v_480P.vae.model } #
offload.profile(pipe, profile_no=4, budgets = {"model":100, "*":3000}, verboseLevel=2, compile="transformer" )
return '480P'
def prompt_enc(prompt, img, tar_lang):
print('prompt extend...')
if img is None:
print('Please upload an image')
return prompt
global prompt_expander
prompt_output = prompt_expander(
prompt, image=img, tar_lang=tar_lang.lower())
if prompt_output.status == False:
return prompt
else:
return prompt_output.prompt
def i2v_generation(img2vid_prompt, img2vid_image, res, sd_steps,
guide_scale, shift_scale, seed, n_prompt):
# print(f"{img2vid_prompt},{resolution},{sd_steps},{guide_scale},{shift_scale},{seed},{n_prompt}")
global resolution
from PIL import Image
img2vid_image = Image.open("d:\mammoth2.jpg")
if resolution == '------':
print(
'Please specify at least one resolution ckpt dir or specify the resolution'
)
return None
else:
if resolution == '720P':
global wan_i2v_720P
video = wan_i2v_720P.generate(
img2vid_prompt,
img2vid_image,
max_area=MAX_AREA_CONFIGS['720*1280'],
shift=shift_scale,
sampling_steps=sd_steps,
guide_scale=guide_scale,
n_prompt=n_prompt,
seed=seed,
offload_model=False)
else:
global wan_i2v_480P
video = wan_i2v_480P.generate(
img2vid_prompt,
img2vid_image,
max_area=MAX_AREA_CONFIGS['480*832'],
shift=3.0, #shift_scale
sampling_steps=sd_steps,
guide_scale=guide_scale,
n_prompt=n_prompt,
seed=seed,
offload_model=False)
cache_video(
tensor=video[None],
save_file="example.mp4",
fps=16,
nrow=1,
normalize=True,
value_range=(-1, 1))
return "example.mp4"
# Interface
def gradio_interface():
with gr.Blocks() as demo:
gr.Markdown("""
<div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
Wan2.1 (I2V-14B)
</div>
<div style="text-align: center; font-size: 16px; font-weight: normal; margin-bottom: 20px;">
Wan: Open and Advanced Large-Scale Video Generative Models.
</div>
""")
with gr.Row():
with gr.Column():
resolution = gr.Dropdown(
label='Resolution',
choices=['------', '720P', '480P'],
value='------')
img2vid_image = gr.Image(
type="pil",
label="Upload Input Image",
elem_id="image_upload",
)
img2vid_prompt = gr.Textbox(
label="Prompt",
value="Several giant wooly mammoths approach treading through a snowy meadow, their long wooly fur lightly blows in the wind as they walk, snow covered trees and dramatic snow capped mountains in the distance, mid afternoon light with wispy clouds and a sun high in the distance creates a warm glow, the low camera view is stunning capturing the large furry mammal with beautiful photography, depth of field.",
placeholder="Describe the video you want to generate",
)
tar_lang = gr.Radio(
choices=["CH", "EN"],
label="Target language of prompt enhance",
value="CH")
run_p_button = gr.Button(value="Prompt Enhance")
with gr.Accordion("Advanced Options", open=True):
with gr.Row():
sd_steps = gr.Slider(
label="Diffusion steps",
minimum=1,
maximum=1000,
value=50,
step=1)
guide_scale = gr.Slider(
label="Guide scale",
minimum=0,
maximum=20,
value=5.0,
step=1)
with gr.Row():
shift_scale = gr.Slider(
label="Shift scale",
minimum=0,
maximum=10,
value=5.0,
step=1)
seed = gr.Slider(
label="Seed",
minimum=-1,
maximum=2147483647,
step=1,
value=-1)
n_prompt = gr.Textbox(
label="Negative Prompt",
placeholder="Describe the negative prompt you want to add"
)
run_i2v_button = gr.Button("Generate Video")
with gr.Column():
result_gallery = gr.Video(
label='Generated Video', interactive=False, height=600)
resolution.input(
fn=load_model, inputs=[resolution], outputs=[resolution])
run_p_button.click(
fn=prompt_enc,
inputs=[img2vid_prompt, img2vid_image, tar_lang],
outputs=[img2vid_prompt])
run_i2v_button.click(
fn=i2v_generation,
inputs=[
img2vid_prompt, img2vid_image, resolution, sd_steps,
guide_scale, shift_scale, seed, n_prompt
],
outputs=[result_gallery],
)
return demo
# Main
def _parse_args():
parser = argparse.ArgumentParser(
description="Generate a video from a text prompt or image using Gradio")
parser.add_argument(
"--ckpt_dir_720p",
type=str,
default=None,
help="The path to the checkpoint directory.")
parser.add_argument(
"--ckpt_dir_480p",
type=str,
default=None,
help="The path to the checkpoint directory.")
parser.add_argument(
"--prompt_extend_method",
type=str,
default="local_qwen",
choices=["dashscope", "local_qwen"],
help="The prompt extend method to use.")
parser.add_argument(
"--prompt_extend_model",
type=str,
default=None,
help="The prompt extend model to use.")
args = parser.parse_args()
args.ckpt_dir_720p = "../ckpts" # os.path.join("ckpt")
args.ckpt_dir_480p = "../ckpts" # os.path.join("ckpt")
assert args.ckpt_dir_720p is not None or args.ckpt_dir_480p is not None, "Please specify at least one checkpoint directory."
return args
if __name__ == '__main__':
args = _parse_args()
global resolution
# load_model('720P')
# resolution = '720P'
resolution = '480P'
load_i2v_model(resolution)
print("Step1: Init prompt_expander...", end='', flush=True)
if args.prompt_extend_method == "dashscope":
prompt_expander = DashScopePromptExpander(
model_name=args.prompt_extend_model, is_vl=True)
elif args.prompt_extend_method == "local_qwen":
prompt_expander = QwenPromptExpander(
model_name=args.prompt_extend_model, is_vl=True, device=0)
else:
raise NotImplementedError(
f"Unsupport prompt_extend_method: {args.prompt_extend_method}")
print("done", flush=True)
demo = gradio_interface()
demo.launch(server_name="0.0.0.0", share=False, server_port=7860)

206
gradio/t2i_14B_singleGPU.py
View File

@@ -1,206 +0,0 @@
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import argparse
import os.path as osp
import os
import sys
import warnings
import gradio as gr
warnings.filterwarnings('ignore')
# Model
sys.path.insert(0, os.path.sep.join(osp.realpath(__file__).split(os.path.sep)[:-2]))
import wan
from wan.configs import WAN_CONFIGS
from wan.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander
from wan.utils.utils import cache_image
# Global Var
prompt_expander = None
wan_t2i = None
# Button Func
def prompt_enc(prompt, tar_lang):
global prompt_expander
prompt_output = prompt_expander(prompt, tar_lang=tar_lang.lower())
if prompt_output.status == False:
return prompt
else:
return prompt_output.prompt
def t2i_generation(txt2img_prompt, resolution, sd_steps, guide_scale,
shift_scale, seed, n_prompt):
global wan_t2i
# print(f"{txt2img_prompt},{resolution},{sd_steps},{guide_scale},{shift_scale},{seed},{n_prompt}")
W = int(resolution.split("*")[0])
H = int(resolution.split("*")[1])
video = wan_t2i.generate(
txt2img_prompt,
size=(W, H),
frame_num=1,
shift=shift_scale,
sampling_steps=sd_steps,
guide_scale=guide_scale,
n_prompt=n_prompt,
seed=seed,
offload_model=True)
cache_image(
tensor=video.squeeze(1)[None],
save_file="example.png",
nrow=1,
normalize=True,
value_range=(-1, 1))
return "example.png"
# Interface
def gradio_interface():
with gr.Blocks() as demo:
gr.Markdown("""
<div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
Wan2.1 (T2I-14B)
</div>
<div style="text-align: center; font-size: 16px; font-weight: normal; margin-bottom: 20px;">
Wan: Open and Advanced Large-Scale Video Generative Models.
</div>
""")
with gr.Row():
with gr.Column():
txt2img_prompt = gr.Textbox(
label="Prompt",
placeholder="Describe the image you want to generate",
)
tar_lang = gr.Radio(
choices=["CH", "EN"],
label="Target language of prompt enhance",
value="CH")
run_p_button = gr.Button(value="Prompt Enhance")
with gr.Accordion("Advanced Options", open=True):
resolution = gr.Dropdown(
label='Resolution(Width*Height)',
choices=[
'720*1280', '1280*720', '960*960', '1088*832',
'832*1088', '480*832', '832*480', '624*624',
'704*544', '544*704'
],
value='720*1280')
with gr.Row():
sd_steps = gr.Slider(
label="Diffusion steps",
minimum=1,
maximum=1000,
value=50,
step=1)
guide_scale = gr.Slider(
label="Guide scale",
minimum=0,
maximum=20,
value=5.0,
step=1)
with gr.Row():
shift_scale = gr.Slider(
label="Shift scale",
minimum=0,
maximum=10,
value=5.0,
step=1)
seed = gr.Slider(
label="Seed",
minimum=-1,
maximum=2147483647,
step=1,
value=-1)
n_prompt = gr.Textbox(
label="Negative Prompt",
placeholder="Describe the negative prompt you want to add"
)
run_t2i_button = gr.Button("Generate Image")
with gr.Column():
result_gallery = gr.Image(
label='Generated Image', interactive=False, height=600)
run_p_button.click(
fn=prompt_enc,
inputs=[txt2img_prompt, tar_lang],
outputs=[txt2img_prompt])
run_t2i_button.click(
fn=t2i_generation,
inputs=[
txt2img_prompt, resolution, sd_steps, guide_scale, shift_scale,
seed, n_prompt
],
outputs=[result_gallery],
)
return demo
# Main
def _parse_args():
parser = argparse.ArgumentParser(
description="Generate a image from a text prompt or image using Gradio")
parser.add_argument(
"--ckpt_dir",
type=str,
default="cache",
help="The path to the checkpoint directory.")
parser.add_argument(
"--prompt_extend_method",
type=str,
default="local_qwen",
choices=["dashscope", "local_qwen"],
help="The prompt extend method to use.")
parser.add_argument(
"--prompt_extend_model",
type=str,
default=None,
help="The prompt extend model to use.")
args = parser.parse_args()
return args
if __name__ == '__main__':
args = _parse_args()
print("Step1: Init prompt_expander...", end='', flush=True)
if args.prompt_extend_method == "dashscope":
prompt_expander = DashScopePromptExpander(
model_name=args.prompt_extend_model, is_vl=False)
elif args.prompt_extend_method == "local_qwen":
prompt_expander = QwenPromptExpander(
model_name=args.prompt_extend_model, is_vl=False, device=0)
else:
raise NotImplementedError(
f"Unsupport prompt_extend_method: {args.prompt_extend_method}")
print("done", flush=True)
print("Step2: Init 14B t2i model...", end='', flush=True)
cfg = WAN_CONFIGS['t2i-14B']
# cfg = WAN_CONFIGS['t2v-1.3B']
wan_t2i = wan.WanT2V(
config=cfg,
checkpoint_dir=args.ckpt_dir,
device_id=0,
rank=0,
t5_fsdp=False,
dit_fsdp=False,
use_usp=False,
)
print("done", flush=True)
demo = gradio_interface()
demo.launch(server_name="0.0.0.0", share=False, server_port=7860)

207
gradio/t2v_1.3B_singleGPU.py
View File

@@ -1,207 +0,0 @@
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import argparse
import os.path as osp
import os
import sys
import warnings
import gradio as gr
warnings.filterwarnings('ignore')
# Model
sys.path.insert(0, os.path.sep.join(osp.realpath(__file__).split(os.path.sep)[:-2]))
import wan
from wan.configs import WAN_CONFIGS
from wan.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander
from wan.utils.utils import cache_video
# Global Var
prompt_expander = None
wan_t2v = None
# Button Func
def prompt_enc(prompt, tar_lang):
global prompt_expander
prompt_output = prompt_expander(prompt, tar_lang=tar_lang.lower())
if prompt_output.status == False:
return prompt
else:
return prompt_output.prompt
def t2v_generation(txt2vid_prompt, resolution, sd_steps, guide_scale,
shift_scale, seed, n_prompt):
global wan_t2v
# print(f"{txt2vid_prompt},{resolution},{sd_steps},{guide_scale},{shift_scale},{seed},{n_prompt}")
W = int(resolution.split("*")[0])
H = int(resolution.split("*")[1])
video = wan_t2v.generate(
txt2vid_prompt,
size=(W, H),
shift=shift_scale,
sampling_steps=sd_steps,
guide_scale=guide_scale,
n_prompt=n_prompt,
seed=seed,
offload_model=True)
cache_video(
tensor=video[None],
save_file="example.mp4",
fps=16,
nrow=1,
normalize=True,
value_range=(-1, 1))
return "example.mp4"
# Interface
def gradio_interface():
with gr.Blocks() as demo:
gr.Markdown("""
<div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
Wan2.1 (T2V-1.3B)
</div>
<div style="text-align: center; font-size: 16px; font-weight: normal; margin-bottom: 20px;">
Wan: Open and Advanced Large-Scale Video Generative Models.
</div>
""")
with gr.Row():
with gr.Column():
txt2vid_prompt = gr.Textbox(
label="Prompt",
placeholder="Describe the video you want to generate",
)
tar_lang = gr.Radio(
choices=["CH", "EN"],
label="Target language of prompt enhance",
value="CH")
run_p_button = gr.Button(value="Prompt Enhance")
with gr.Accordion("Advanced Options", open=True):
resolution = gr.Dropdown(
label='Resolution(Width*Height)',
choices=[
'480*832',
'832*480',
'624*624',
'704*544',
'544*704',
],
value='480*832')
with gr.Row():
sd_steps = gr.Slider(
label="Diffusion steps",
minimum=1,
maximum=1000,
value=50,
step=1)
guide_scale = gr.Slider(
label="Guide scale",
minimum=0,
maximum=20,
value=6.0,
step=1)
with gr.Row():
shift_scale = gr.Slider(
label="Shift scale",
minimum=0,
maximum=20,
value=8.0,
step=1)
seed = gr.Slider(
label="Seed",
minimum=-1,
maximum=2147483647,
step=1,
value=-1)
n_prompt = gr.Textbox(
label="Negative Prompt",
placeholder="Describe the negative prompt you want to add"
)
run_t2v_button = gr.Button("Generate Video")
with gr.Column():
result_gallery = gr.Video(
label='Generated Video', interactive=False, height=600)
run_p_button.click(
fn=prompt_enc,
inputs=[txt2vid_prompt, tar_lang],
outputs=[txt2vid_prompt])
run_t2v_button.click(
fn=t2v_generation,
inputs=[
txt2vid_prompt, resolution, sd_steps, guide_scale, shift_scale,
seed, n_prompt
],
outputs=[result_gallery],
)
return demo
# Main
def _parse_args():
parser = argparse.ArgumentParser(
description="Generate a video from a text prompt or image using Gradio")
parser.add_argument(
"--ckpt_dir",
type=str,
default="cache",
help="The path to the checkpoint directory.")
parser.add_argument(
"--prompt_extend_method",
type=str,
default="local_qwen",
choices=["dashscope", "local_qwen"],
help="The prompt extend method to use.")
parser.add_argument(
"--prompt_extend_model",
type=str,
default=None,
help="The prompt extend model to use.")
args = parser.parse_args()
return args
if __name__ == '__main__':
args = _parse_args()
print("Step1: Init prompt_expander...", end='', flush=True)
if args.prompt_extend_method == "dashscope":
prompt_expander = DashScopePromptExpander(
model_name=args.prompt_extend_model, is_vl=False)
elif args.prompt_extend_method == "local_qwen":
prompt_expander = QwenPromptExpander(
model_name=args.prompt_extend_model, is_vl=False, device=0)
else:
raise NotImplementedError(
f"Unsupport prompt_extend_method: {args.prompt_extend_method}")
print("done", flush=True)
print("Step2: Init 1.3B t2v model...", end='', flush=True)
cfg = WAN_CONFIGS['t2v-1.3B']
wan_t2v = wan.WanT2V(
config=cfg,
checkpoint_dir=args.ckpt_dir,
device_id=0,
rank=0,
t5_fsdp=False,
dit_fsdp=False,
use_usp=False,
)
print("done", flush=True)
demo = gradio_interface()
demo.launch(server_name="0.0.0.0", share=False, server_port=7860)

216
gradio/t2v_14B_singleGPU.py
View File

@@ -1,216 +0,0 @@
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import argparse
import os.path as osp
import os
import sys
import warnings
import gradio as gr
warnings.filterwarnings('ignore')
# Model
sys.path.insert(0, os.path.sep.join(osp.realpath(__file__).split(os.path.sep)[:-2]))
import wan
from wan.configs import WAN_CONFIGS
from wan.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander
from wan.utils.utils import cache_video
# Global Var
prompt_expander = None
wan_t2v = None
# Button Func
def prompt_enc(prompt, tar_lang):
global prompt_expander
prompt_output = prompt_expander(prompt, tar_lang=tar_lang.lower())
if prompt_output.status == False:
return prompt
else:
return prompt_output.prompt
def t2v_generation(txt2vid_prompt, resolution, sd_steps, guide_scale,
shift_scale, seed, n_prompt):
global wan_t2v
# print(f"{txt2vid_prompt},{resolution},{sd_steps},{guide_scale},{shift_scale},{seed},{n_prompt}")
W = int(resolution.split("*")[0])
H = int(resolution.split("*")[1])
video = wan_t2v.generate(
txt2vid_prompt,
size=(W, H),
shift=shift_scale,
sampling_steps=sd_steps,
guide_scale=guide_scale,
n_prompt=n_prompt,
seed=seed,
offload_model=False)
cache_video(
tensor=video[None],
save_file="example.mp4",
fps=16,
nrow=1,
normalize=True,
value_range=(-1, 1))
return "example.mp4"
# Interface
def gradio_interface():
with gr.Blocks() as demo:
gr.Markdown("""
<div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
Wan2.1 (T2V-14B)
</div>
<div style="text-align: center; font-size: 16px; font-weight: normal; margin-bottom: 20px;">
Wan: Open and Advanced Large-Scale Video Generative Models.
</div>
""")
with gr.Row():
with gr.Column():
txt2vid_prompt = gr.Textbox(
label="Prompt",
placeholder="Describe the video you want to generate",
)
tar_lang = gr.Radio(
choices=["CH", "EN"],
label="Target language of prompt enhance",
value="CH")
run_p_button = gr.Button(value="Prompt Enhance")
with gr.Accordion("Advanced Options", open=True):
resolution = gr.Dropdown(
label='Resolution(Width*Height)',
choices=[
'720*1280', '1280*720', '960*960', '1088*832',
'832*1088', '480*832', '832*480', '624*624',
'704*544', '544*704'
],
value='720*1280')
with gr.Row():
sd_steps = gr.Slider(
label="Diffusion steps",
minimum=1,
maximum=1000,
value=50,
step=1)
guide_scale = gr.Slider(
label="Guide scale",
minimum=0,
maximum=20,
value=5.0,
step=1)
with gr.Row():
shift_scale = gr.Slider(
label="Shift scale",
minimum=0,
maximum=10,
value=5.0,
step=1)
seed = gr.Slider(
label="Seed",
minimum=-1,
maximum=2147483647,
step=1,
value=-1)
n_prompt = gr.Textbox(
label="Negative Prompt",
placeholder="Describe the negative prompt you want to add"
)
run_t2v_button = gr.Button("Generate Video")
with gr.Column():
result_gallery = gr.Video(
label='Generated Video', interactive=False, height=600)
run_p_button.click(
fn=prompt_enc,
inputs=[txt2vid_prompt, tar_lang],
outputs=[txt2vid_prompt])
run_t2v_button.click(
fn=t2v_generation,
inputs=[
txt2vid_prompt, resolution, sd_steps, guide_scale, shift_scale,
seed, n_prompt
],
outputs=[result_gallery],
)
return demo
# Main
def _parse_args():
parser = argparse.ArgumentParser(
description="Generate a video from a text prompt or image using Gradio")
parser.add_argument(
"--ckpt_dir",
type=str,
default="cache",
help="The path to the checkpoint directory.")
parser.add_argument(
"--prompt_extend_method",
type=str,
default="local_qwen",
choices=["dashscope", "local_qwen"],
help="The prompt extend method to use.")
parser.add_argument(
"--prompt_extend_model",
type=str,
default=None,
help="The prompt extend model to use.")
args = parser.parse_args()
return args
if __name__ == '__main__':
args = _parse_args()
print("Step1: Init prompt_expander...", end='', flush=True)
prompt_expander = None
# if args.prompt_extend_method == "dashscope":
# prompt_expander = DashScopePromptExpander(
# model_name=args.prompt_extend_model, is_vl=False)
# elif args.prompt_extend_method == "local_qwen":
# prompt_expander = QwenPromptExpander(
# model_name=args.prompt_extend_model, is_vl=False, device=0)
# else:
# raise NotImplementedError(
# f"Unsupport prompt_extend_method: {args.prompt_extend_method}")
# print("done", flush=True)
from mmgp import offload
print("Step2: Init 14B t2v model...", end='', flush=True)
cfg = WAN_CONFIGS['t2v-14B']
# cfg = WAN_CONFIGS['t2v-1.3B']
wan_t2v = wan.WanT2V(
config=cfg,
checkpoint_dir="../ckpts",
device_id=0,
rank=0,
t5_fsdp=False,
dit_fsdp=False,
use_usp=False,
)
pipe = {"transformer": wan_t2v.model, "text_encoder" : wan_t2v.text_encoder.model, "vae": wan_t2v.vae.model } #
# offload.profile(pipe, profile_no=4, budgets = {"transformer":100, "*":3000}, verboseLevel=2, quantizeTransformer = False, compile = "transformer") #
offload.profile(pipe, profile_no=4, budgets = {"transformer":100, "*":3000}, verboseLevel=2, quantizeTransformer = False) #
# offload.profile(pipe, profile_no=4, budgets = {"transformer":3000, "*":3000}, verboseLevel=2, quantizeTransformer = False)
print("done", flush=True)
demo = gradio_interface()
demo.launch(server_name="0.0.0.0", share=False, server_port=7860)

2457
gradio_server.py
View File

File diff suppressed because it is too large Load Diff

2
preprocessing/dwpose/__init__.py Normal file
View File

@@ -0,0 +1,2 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.

127
preprocessing/dwpose/onnxdet.py Normal file
View File

@@ -0,0 +1,127 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import cv2
import numpy as np
import onnxruntime
def nms(boxes, scores, nms_thr):
"""Single class NMS implemented in Numpy."""
x1 = boxes[:, 0]
y1 = boxes[:, 1]
x2 = boxes[:, 2]
y2 = boxes[:, 3]
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
order = scores.argsort()[::-1]
keep = []
while order.size > 0:
i = order[0]
keep.append(i)
xx1 = np.maximum(x1[i], x1[order[1:]])
yy1 = np.maximum(y1[i], y1[order[1:]])
xx2 = np.minimum(x2[i], x2[order[1:]])
yy2 = np.minimum(y2[i], y2[order[1:]])
w = np.maximum(0.0, xx2 - xx1 + 1)
h = np.maximum(0.0, yy2 - yy1 + 1)
inter = w * h
ovr = inter / (areas[i] + areas[order[1:]] - inter)
inds = np.where(ovr <= nms_thr)[0]
order = order[inds + 1]
return keep
def multiclass_nms(boxes, scores, nms_thr, score_thr):
"""Multiclass NMS implemented in Numpy. Class-aware version."""
final_dets = []
num_classes = scores.shape[1]
for cls_ind in range(num_classes):
cls_scores = scores[:, cls_ind]
valid_score_mask = cls_scores > score_thr
if valid_score_mask.sum() == 0:
continue
else:
valid_scores = cls_scores[valid_score_mask]
valid_boxes = boxes[valid_score_mask]
keep = nms(valid_boxes, valid_scores, nms_thr)
if len(keep) > 0:
cls_inds = np.ones((len(keep), 1)) * cls_ind
dets = np.concatenate(
[valid_boxes[keep], valid_scores[keep, None], cls_inds], 1
)
final_dets.append(dets)
if len(final_dets) == 0:
return None
return np.concatenate(final_dets, 0)
def demo_postprocess(outputs, img_size, p6=False):
grids = []
expanded_strides = []
strides = [8, 16, 32] if not p6 else [8, 16, 32, 64]
hsizes = [img_size[0] // stride for stride in strides]
wsizes = [img_size[1] // stride for stride in strides]
for hsize, wsize, stride in zip(hsizes, wsizes, strides):
xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize))
grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
grids.append(grid)
shape = grid.shape[:2]
expanded_strides.append(np.full((*shape, 1), stride))
grids = np.concatenate(grids, 1)
expanded_strides = np.concatenate(expanded_strides, 1)
outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides
return outputs
def preprocess(img, input_size, swap=(2, 0, 1)):
if len(img.shape) == 3:
padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114
else:
padded_img = np.ones(input_size, dtype=np.uint8) * 114
r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
resized_img = cv2.resize(
img,
(int(img.shape[1] * r), int(img.shape[0] * r)),
interpolation=cv2.INTER_LINEAR,
).astype(np.uint8)
padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img
padded_img = padded_img.transpose(swap)
padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
return padded_img, r
def inference_detector(session, oriImg):
input_shape = (640,640)
img, ratio = preprocess(oriImg, input_shape)
ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
output = session.run(None, ort_inputs)
predictions = demo_postprocess(output[0], input_shape)[0]
boxes = predictions[:, :4]
scores = predictions[:, 4:5] * predictions[:, 5:]
boxes_xyxy = np.ones_like(boxes)
boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
boxes_xyxy /= ratio
dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)
if dets is not None:
final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
isscore = final_scores>0.3
iscat = final_cls_inds == 0
isbbox = [ i and j for (i, j) in zip(isscore, iscat)]
final_boxes = final_boxes[isbbox]
else:
final_boxes = np.array([])
return final_boxes

362
preprocessing/dwpose/onnxpose.py Normal file
View File

@@ -0,0 +1,362 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
from typing import List, Tuple
import cv2
import numpy as np
import onnxruntime as ort
def preprocess(
img: np.ndarray, out_bbox, input_size: Tuple[int, int] = (192, 256)
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Do preprocessing for RTMPose model inference.
Args:
img (np.ndarray): Input image in shape.
input_size (tuple): Input image size in shape (w, h).
Returns:
tuple:
- resized_img (np.ndarray): Preprocessed image.
- center (np.ndarray): Center of image.
- scale (np.ndarray): Scale of image.
"""
# get shape of image
img_shape = img.shape[:2]
out_img, out_center, out_scale = [], [], []
if len(out_bbox) == 0:
out_bbox = [[0, 0, img_shape[1], img_shape[0]]]
for i in range(len(out_bbox)):
x0 = out_bbox[i][0]
y0 = out_bbox[i][1]
x1 = out_bbox[i][2]
y1 = out_bbox[i][3]
bbox = np.array([x0, y0, x1, y1])
# get center and scale
center, scale = bbox_xyxy2cs(bbox, padding=1.25)
# do affine transformation
resized_img, scale = top_down_affine(input_size, scale, center, img)
# normalize image
mean = np.array([123.675, 116.28, 103.53])
std = np.array([58.395, 57.12, 57.375])
resized_img = (resized_img - mean) / std
out_img.append(resized_img)
out_center.append(center)
out_scale.append(scale)
return out_img, out_center, out_scale
def inference(sess: ort.InferenceSession, img: np.ndarray) -> np.ndarray:
"""Inference RTMPose model.
Args:
sess (ort.InferenceSession): ONNXRuntime session.
img (np.ndarray): Input image in shape.
Returns:
outputs (np.ndarray): Output of RTMPose model.
"""
all_out = []
# build input
for i in range(len(img)):
input = [img[i].transpose(2, 0, 1)]
# build output
sess_input = {sess.get_inputs()[0].name: input}
sess_output = []
for out in sess.get_outputs():
sess_output.append(out.name)
# run model
outputs = sess.run(sess_output, sess_input)
all_out.append(outputs)
return all_out
def postprocess(outputs: List[np.ndarray],
model_input_size: Tuple[int, int],
center: Tuple[int, int],
scale: Tuple[int, int],
simcc_split_ratio: float = 2.0
) -> Tuple[np.ndarray, np.ndarray]:
"""Postprocess for RTMPose model output.
Args:
outputs (np.ndarray): Output of RTMPose model.
model_input_size (tuple): RTMPose model Input image size.
center (tuple): Center of bbox in shape (x, y).
scale (tuple): Scale of bbox in shape (w, h).
simcc_split_ratio (float): Split ratio of simcc.
Returns:
tuple:
- keypoints (np.ndarray): Rescaled keypoints.
- scores (np.ndarray): Model predict scores.
"""
all_key = []
all_score = []
for i in range(len(outputs)):
# use simcc to decode
simcc_x, simcc_y = outputs[i]
keypoints, scores = decode(simcc_x, simcc_y, simcc_split_ratio)
# rescale keypoints
keypoints = keypoints / model_input_size * scale[i] + center[i] - scale[i] / 2
all_key.append(keypoints[0])
all_score.append(scores[0])
return np.array(all_key), np.array(all_score)
def bbox_xyxy2cs(bbox: np.ndarray,
padding: float = 1.) -> Tuple[np.ndarray, np.ndarray]:
"""Transform the bbox format from (x,y,w,h) into (center, scale)
Args:
bbox (ndarray): Bounding box(es) in shape (4,) or (n, 4), formatted
as (left, top, right, bottom)
padding (float): BBox padding factor that will be multilied to scale.
Default: 1.0
Returns:
tuple: A tuple containing center and scale.
- np.ndarray[float32]: Center (x, y) of the bbox in shape (2,) or
(n, 2)
- np.ndarray[float32]: Scale (w, h) of the bbox in shape (2,) or
(n, 2)
"""
# convert single bbox from (4, ) to (1, 4)
dim = bbox.ndim
if dim == 1:
bbox = bbox[None, :]
# get bbox center and scale
x1, y1, x2, y2 = np.hsplit(bbox, [1, 2, 3])
center = np.hstack([x1 + x2, y1 + y2]) * 0.5
scale = np.hstack([x2 - x1, y2 - y1]) * padding
if dim == 1:
center = center[0]
scale = scale[0]
return center, scale
def _fix_aspect_ratio(bbox_scale: np.ndarray,
aspect_ratio: float) -> np.ndarray:
"""Extend the scale to match the given aspect ratio.
Args:
scale (np.ndarray): The image scale (w, h) in shape (2, )
aspect_ratio (float): The ratio of ``w/h``
Returns:
np.ndarray: The reshaped image scale in (2, )
"""
w, h = np.hsplit(bbox_scale, [1])
bbox_scale = np.where(w > h * aspect_ratio,
np.hstack([w, w / aspect_ratio]),
np.hstack([h * aspect_ratio, h]))
return bbox_scale
def _rotate_point(pt: np.ndarray, angle_rad: float) -> np.ndarray:
"""Rotate a point by an angle.
Args:
pt (np.ndarray): 2D point coordinates (x, y) in shape (2, )
angle_rad (float): rotation angle in radian
Returns:
np.ndarray: Rotated point in shape (2, )
"""
sn, cs = np.sin(angle_rad), np.cos(angle_rad)
rot_mat = np.array([[cs, -sn], [sn, cs]])
return rot_mat @ pt
def _get_3rd_point(a: np.ndarray, b: np.ndarray) -> np.ndarray:
"""To calculate the affine matrix, three pairs of points are required. This
function is used to get the 3rd point, given 2D points a & b.
The 3rd point is defined by rotating vector `a - b` by 90 degrees
anticlockwise, using b as the rotation center.
Args:
a (np.ndarray): The 1st point (x,y) in shape (2, )
b (np.ndarray): The 2nd point (x,y) in shape (2, )
Returns:
np.ndarray: The 3rd point.
"""
direction = a - b
c = b + np.r_[-direction[1], direction[0]]
return c
def get_warp_matrix(center: np.ndarray,
scale: np.ndarray,
rot: float,
output_size: Tuple[int, int],
shift: Tuple[float, float] = (0., 0.),
inv: bool = False) -> np.ndarray:
"""Calculate the affine transformation matrix that can warp the bbox area
in the input image to the output size.
Args:
center (np.ndarray[2, ]): Center of the bounding box (x, y).
scale (np.ndarray[2, ]): Scale of the bounding box
wrt [width, height].
rot (float): Rotation angle (degree).
output_size (np.ndarray[2, ] | list(2,)): Size of the
destination heatmaps.
shift (0-100%): Shift translation ratio wrt the width/height.
Default (0., 0.).
inv (bool): Option to inverse the affine transform direction.
(inv=False: src->dst or inv=True: dst->src)
Returns:
np.ndarray: A 2x3 transformation matrix
"""
shift = np.array(shift)
src_w = scale[0]
dst_w = output_size[0]
dst_h = output_size[1]
# compute transformation matrix
rot_rad = np.deg2rad(rot)
src_dir = _rotate_point(np.array([0., src_w * -0.5]), rot_rad)
dst_dir = np.array([0., dst_w * -0.5])
# get four corners of the src rectangle in the original image
src = np.zeros((3, 2), dtype=np.float32)
src[0, :] = center + scale * shift
src[1, :] = center + src_dir + scale * shift
src[2, :] = _get_3rd_point(src[0, :], src[1, :])
# get four corners of the dst rectangle in the input image
dst = np.zeros((3, 2), dtype=np.float32)
dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :])
if inv:
warp_mat = cv2.getAffineTransform(np.float32(dst), np.float32(src))
else:
warp_mat = cv2.getAffineTransform(np.float32(src), np.float32(dst))
return warp_mat
def top_down_affine(input_size: dict, bbox_scale: dict, bbox_center: dict,
img: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
"""Get the bbox image as the model input by affine transform.
Args:
input_size (dict): The input size of the model.
bbox_scale (dict): The bbox scale of the img.
bbox_center (dict): The bbox center of the img.
img (np.ndarray): The original image.
Returns:
tuple: A tuple containing center and scale.
- np.ndarray[float32]: img after affine transform.
- np.ndarray[float32]: bbox scale after affine transform.
"""
w, h = input_size
warp_size = (int(w), int(h))
# reshape bbox to fixed aspect ratio
bbox_scale = _fix_aspect_ratio(bbox_scale, aspect_ratio=w / h)
# get the affine matrix
center = bbox_center
scale = bbox_scale
rot = 0
warp_mat = get_warp_matrix(center, scale, rot, output_size=(w, h))
# do affine transform
img = cv2.warpAffine(img, warp_mat, warp_size, flags=cv2.INTER_LINEAR)
return img, bbox_scale
def get_simcc_maximum(simcc_x: np.ndarray,
simcc_y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
"""Get maximum response location and value from simcc representations.
Note:
instance number: N
num_keypoints: K
heatmap height: H
heatmap width: W
Args:
simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx)
simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy)
Returns:
tuple:
- locs (np.ndarray): locations of maximum heatmap responses in shape
(K, 2) or (N, K, 2)
- vals (np.ndarray): values of maximum heatmap responses in shape
(K,) or (N, K)
"""
N, K, Wx = simcc_x.shape
simcc_x = simcc_x.reshape(N * K, -1)
simcc_y = simcc_y.reshape(N * K, -1)
# get maximum value locations
x_locs = np.argmax(simcc_x, axis=1)
y_locs = np.argmax(simcc_y, axis=1)
locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)
max_val_x = np.amax(simcc_x, axis=1)
max_val_y = np.amax(simcc_y, axis=1)
# get maximum value across x and y axis
mask = max_val_x > max_val_y
max_val_x[mask] = max_val_y[mask]
vals = max_val_x
locs[vals <= 0.] = -1
# reshape
locs = locs.reshape(N, K, 2)
vals = vals.reshape(N, K)
return locs, vals
def decode(simcc_x: np.ndarray, simcc_y: np.ndarray,
simcc_split_ratio) -> Tuple[np.ndarray, np.ndarray]:
"""Modulate simcc distribution with Gaussian.
Args:
simcc_x (np.ndarray[K, Wx]): model predicted simcc in x.
simcc_y (np.ndarray[K, Wy]): model predicted simcc in y.
simcc_split_ratio (int): The split ratio of simcc.
Returns:
tuple: A tuple containing center and scale.
- np.ndarray[float32]: keypoints in shape (K, 2) or (n, K, 2)
- np.ndarray[float32]: scores in shape (K,) or (n, K)
"""
keypoints, scores = get_simcc_maximum(simcc_x, simcc_y)
keypoints /= simcc_split_ratio
return keypoints, scores
def inference_pose(session, out_bbox, oriImg):
h, w = session.get_inputs()[0].shape[2:]
model_input_size = (w, h)
resized_img, center, scale = preprocess(oriImg, out_bbox, model_input_size)
outputs = inference(session, resized_img)
keypoints, scores = postprocess(outputs, model_input_size, center, scale)
return keypoints, scores

183
preprocessing/dwpose/pose.py Normal file
View File

@@ -0,0 +1,183 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import os
import cv2
import torch
import numpy as np
from . import util
from .wholebody import Wholebody, HWC3, resize_image
from PIL import Image
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
def convert_to_numpy(image):
if isinstance(image, Image.Image):
image = np.array(image)
elif isinstance(image, torch.Tensor):
image = image.detach().cpu().numpy()
elif isinstance(image, np.ndarray):
image = image.copy()
else:
raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
return image
def draw_pose(pose, H, W, use_hand=False, use_body=False, use_face=False):
bodies = pose['bodies']
faces = pose['faces']
hands = pose['hands']
candidate = bodies['candidate']
subset = bodies['subset']
canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
if use_body:
canvas = util.draw_bodypose(canvas, candidate, subset)
if use_hand:
canvas = util.draw_handpose(canvas, hands)
if use_face:
canvas = util.draw_facepose(canvas, faces)
return canvas
class PoseAnnotator:
def __init__(self, cfg, device=None):
onnx_det = cfg['DETECTION_MODEL']
onnx_pose = cfg['POSE_MODEL']
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if device is None else device
self.pose_estimation = Wholebody(onnx_det, onnx_pose, device=self.device)
self.resize_size = cfg.get("RESIZE_SIZE", 1024)
self.use_body = cfg.get('USE_BODY', True)
self.use_face = cfg.get('USE_FACE', True)
self.use_hand = cfg.get('USE_HAND', True)
@torch.no_grad()
@torch.inference_mode
def forward(self, image):
image = convert_to_numpy(image)
input_image = HWC3(image[..., ::-1])
return self.process(resize_image(input_image, self.resize_size), image.shape[:2])
def process(self, ori_img, ori_shape):
ori_h, ori_w = ori_shape
ori_img = ori_img.copy()
H, W, C = ori_img.shape
with torch.no_grad():
candidate, subset, det_result = self.pose_estimation(ori_img)
nums, keys, locs = candidate.shape
candidate[..., 0] /= float(W)
candidate[..., 1] /= float(H)
body = candidate[:, :18].copy()
body = body.reshape(nums * 18, locs)
score = subset[:, :18]
for i in range(len(score)):
for j in range(len(score[i])):
if score[i][j] > 0.3:
score[i][j] = int(18 * i + j)
else:
score[i][j] = -1
un_visible = subset < 0.3
candidate[un_visible] = -1
foot = candidate[:, 18:24]
faces = candidate[:, 24:92]
hands = candidate[:, 92:113]
hands = np.vstack([hands, candidate[:, 113:]])
bodies = dict(candidate=body, subset=score)
pose = dict(bodies=bodies, hands=hands, faces=faces)
ret_data = {}
if self.use_body:
detected_map_body = draw_pose(pose, H, W, use_body=True)
detected_map_body = cv2.resize(detected_map_body[..., ::-1], (ori_w, ori_h),
interpolation=cv2.INTER_LANCZOS4 if ori_h * ori_w > H * W else cv2.INTER_AREA)
ret_data["detected_map_body"] = detected_map_body
if self.use_face:
detected_map_face = draw_pose(pose, H, W, use_face=True)
detected_map_face = cv2.resize(detected_map_face[..., ::-1], (ori_w, ori_h),
interpolation=cv2.INTER_LANCZOS4 if ori_h * ori_w > H * W else cv2.INTER_AREA)
ret_data["detected_map_face"] = detected_map_face
if self.use_body and self.use_face:
detected_map_bodyface = draw_pose(pose, H, W, use_body=True, use_face=True)
detected_map_bodyface = cv2.resize(detected_map_bodyface[..., ::-1], (ori_w, ori_h),
interpolation=cv2.INTER_LANCZOS4 if ori_h * ori_w > H * W else cv2.INTER_AREA)
ret_data["detected_map_bodyface"] = detected_map_bodyface
if self.use_hand and self.use_body and self.use_face:
detected_map_handbodyface = draw_pose(pose, H, W, use_hand=True, use_body=True, use_face=True)
detected_map_handbodyface = cv2.resize(detected_map_handbodyface[..., ::-1], (ori_w, ori_h),
interpolation=cv2.INTER_LANCZOS4 if ori_h * ori_w > H * W else cv2.INTER_AREA)
ret_data["detected_map_handbodyface"] = detected_map_handbodyface
# convert_size
if det_result.shape[0] > 0:
w_ratio, h_ratio = ori_w / W, ori_h / H
det_result[..., ::2] *= h_ratio
det_result[..., 1::2] *= w_ratio
det_result = det_result.astype(np.int32)
return ret_data, det_result
class PoseBodyFaceAnnotator(PoseAnnotator):
def __init__(self, cfg):
super().__init__(cfg)
self.use_body, self.use_face, self.use_hand = True, True, False
@torch.no_grad()
@torch.inference_mode
def forward(self, image):
ret_data, det_result = super().forward(image)
return ret_data['detected_map_bodyface']
class PoseBodyFaceVideoAnnotator(PoseBodyFaceAnnotator):
def forward(self, frames):
ret_frames = []
for frame in frames:
anno_frame = super().forward(np.array(frame))
ret_frames.append(anno_frame)
return ret_frames
import imageio
def save_one_video(file_path, videos, fps=8, quality=8, macro_block_size=None):
try:
video_writer = imageio.get_writer(file_path, fps=fps, codec='libx264', quality=quality, macro_block_size=macro_block_size)
for frame in videos:
video_writer.append_data(frame)
video_writer.close()
return True
except Exception as e:
print(f"Video save error: {e}")
return False
def get_frames(video_path):
frames = []
# Opens the Video file with CV2
cap = cv2.VideoCapture(video_path)
fps = cap.get(cv2.CAP_PROP_FPS)
print("video fps: " + str(fps))
i = 0
while cap.isOpened():
ret, frame = cap.read()
if ret == False:
break
frames.append(frame)
i += 1
cap.release()
cv2.destroyAllWindows()
return frames, fps

299
preprocessing/dwpose/util.py Normal file
View File

@@ -0,0 +1,299 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import math
import numpy as np
import matplotlib
import cv2
eps = 0.01
def smart_resize(x, s):
Ht, Wt = s
if x.ndim == 2:
Ho, Wo = x.shape
Co = 1
else:
Ho, Wo, Co = x.shape
if Co == 3 or Co == 1:
k = float(Ht + Wt) / float(Ho + Wo)
return cv2.resize(x, (int(Wt), int(Ht)), interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4)
else:
return np.stack([smart_resize(x[:, :, i], s) for i in range(Co)], axis=2)
def smart_resize_k(x, fx, fy):
if x.ndim == 2:
Ho, Wo = x.shape
Co = 1
else:
Ho, Wo, Co = x.shape
Ht, Wt = Ho * fy, Wo * fx
if Co == 3 or Co == 1:
k = float(Ht + Wt) / float(Ho + Wo)
return cv2.resize(x, (int(Wt), int(Ht)), interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4)
else:
return np.stack([smart_resize_k(x[:, :, i], fx, fy) for i in range(Co)], axis=2)
def padRightDownCorner(img, stride, padValue):
h = img.shape[0]
w = img.shape[1]
pad = 4 * [None]
pad[0] = 0 # up
pad[1] = 0 # left
pad[2] = 0 if (h % stride == 0) else stride - (h % stride) # down
pad[3] = 0 if (w % stride == 0) else stride - (w % stride) # right
img_padded = img
pad_up = np.tile(img_padded[0:1, :, :]*0 + padValue, (pad[0], 1, 1))
img_padded = np.concatenate((pad_up, img_padded), axis=0)
pad_left = np.tile(img_padded[:, 0:1, :]*0 + padValue, (1, pad[1], 1))
img_padded = np.concatenate((pad_left, img_padded), axis=1)
pad_down = np.tile(img_padded[-2:-1, :, :]*0 + padValue, (pad[2], 1, 1))
img_padded = np.concatenate((img_padded, pad_down), axis=0)
pad_right = np.tile(img_padded[:, -2:-1, :]*0 + padValue, (1, pad[3], 1))
img_padded = np.concatenate((img_padded, pad_right), axis=1)
return img_padded, pad
def transfer(model, model_weights):
transfered_model_weights = {}
for weights_name in model.state_dict().keys():
transfered_model_weights[weights_name] = model_weights['.'.join(weights_name.split('.')[1:])]
return transfered_model_weights
def draw_bodypose(canvas, candidate, subset):
H, W, C = canvas.shape
candidate = np.array(candidate)
subset = np.array(subset)
stickwidth = 4
limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \
[10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \
[1, 16], [16, 18], [3, 17], [6, 18]]
colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
[0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
[170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
for i in range(17):
for n in range(len(subset)):
index = subset[n][np.array(limbSeq[i]) - 1]
if -1 in index:
continue
Y = candidate[index.astype(int), 0] * float(W)
X = candidate[index.astype(int), 1] * float(H)
mX = np.mean(X)
mY = np.mean(Y)
length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
cv2.fillConvexPoly(canvas, polygon, colors[i])
canvas = (canvas * 0.6).astype(np.uint8)
for i in range(18):
for n in range(len(subset)):
index = int(subset[n][i])
if index == -1:
continue
x, y = candidate[index][0:2]
x = int(x * W)
y = int(y * H)
cv2.circle(canvas, (int(x), int(y)), 4, colors[i], thickness=-1)
return canvas
def draw_handpose(canvas, all_hand_peaks):
H, W, C = canvas.shape
edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \
[10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]]
for peaks in all_hand_peaks:
peaks = np.array(peaks)
for ie, e in enumerate(edges):
x1, y1 = peaks[e[0]]
x2, y2 = peaks[e[1]]
x1 = int(x1 * W)
y1 = int(y1 * H)
x2 = int(x2 * W)
y2 = int(y2 * H)
if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
cv2.line(canvas, (x1, y1), (x2, y2), matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * 255, thickness=2)
for i, keyponit in enumerate(peaks):
x, y = keyponit
x = int(x * W)
y = int(y * H)
if x > eps and y > eps:
cv2.circle(canvas, (x, y), 4, (0, 0, 255), thickness=-1)
return canvas
def draw_facepose(canvas, all_lmks):
H, W, C = canvas.shape
for lmks in all_lmks:
lmks = np.array(lmks)
for lmk in lmks:
x, y = lmk
x = int(x * W)
y = int(y * H)
if x > eps and y > eps:
cv2.circle(canvas, (x, y), 3, (255, 255, 255), thickness=-1)
return canvas
# detect hand according to body pose keypoints
# please refer to https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/src/openpose/hand/handDetector.cpp
def handDetect(candidate, subset, oriImg):
# right hand: wrist 4, elbow 3, shoulder 2
# left hand: wrist 7, elbow 6, shoulder 5
ratioWristElbow = 0.33
detect_result = []
image_height, image_width = oriImg.shape[0:2]
for person in subset.astype(int):
# if any of three not detected
has_left = np.sum(person[[5, 6, 7]] == -1) == 0
has_right = np.sum(person[[2, 3, 4]] == -1) == 0
if not (has_left or has_right):
continue
hands = []
#left hand
if has_left:
left_shoulder_index, left_elbow_index, left_wrist_index = person[[5, 6, 7]]
x1, y1 = candidate[left_shoulder_index][:2]
x2, y2 = candidate[left_elbow_index][:2]
x3, y3 = candidate[left_wrist_index][:2]
hands.append([x1, y1, x2, y2, x3, y3, True])
# right hand
if has_right:
right_shoulder_index, right_elbow_index, right_wrist_index = person[[2, 3, 4]]
x1, y1 = candidate[right_shoulder_index][:2]
x2, y2 = candidate[right_elbow_index][:2]
x3, y3 = candidate[right_wrist_index][:2]
hands.append([x1, y1, x2, y2, x3, y3, False])
for x1, y1, x2, y2, x3, y3, is_left in hands:
# pos_hand = pos_wrist + ratio * (pos_wrist - pos_elbox) = (1 + ratio) * pos_wrist - ratio * pos_elbox
# handRectangle.x = posePtr[wrist*3] + ratioWristElbow * (posePtr[wrist*3] - posePtr[elbow*3]);
# handRectangle.y = posePtr[wrist*3+1] + ratioWristElbow * (posePtr[wrist*3+1] - posePtr[elbow*3+1]);
# const auto distanceWristElbow = getDistance(poseKeypoints, person, wrist, elbow);
# const auto distanceElbowShoulder = getDistance(poseKeypoints, person, elbow, shoulder);
# handRectangle.width = 1.5f * fastMax(distanceWristElbow, 0.9f * distanceElbowShoulder);
x = x3 + ratioWristElbow * (x3 - x2)
y = y3 + ratioWristElbow * (y3 - y2)
distanceWristElbow = math.sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2)
distanceElbowShoulder = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
width = 1.5 * max(distanceWristElbow, 0.9 * distanceElbowShoulder)
# x-y refers to the center --> offset to topLeft point
# handRectangle.x -= handRectangle.width / 2.f;
# handRectangle.y -= handRectangle.height / 2.f;
x -= width / 2
y -= width / 2 # width = height
# overflow the image
if x < 0: x = 0
if y < 0: y = 0
width1 = width
width2 = width
if x + width > image_width: width1 = image_width - x
if y + width > image_height: width2 = image_height - y
width = min(width1, width2)
# the max hand box value is 20 pixels
if width >= 20:
detect_result.append([int(x), int(y), int(width), is_left])
'''
return value: [[x, y, w, True if left hand else False]].
width=height since the network require squared input.
x, y is the coordinate of top left
'''
return detect_result
# Written by Lvmin
def faceDetect(candidate, subset, oriImg):
# left right eye ear 14 15 16 17
detect_result = []
image_height, image_width = oriImg.shape[0:2]
for person in subset.astype(int):
has_head = person[0] > -1
if not has_head:
continue
has_left_eye = person[14] > -1
has_right_eye = person[15] > -1
has_left_ear = person[16] > -1
has_right_ear = person[17] > -1
if not (has_left_eye or has_right_eye or has_left_ear or has_right_ear):
continue
head, left_eye, right_eye, left_ear, right_ear = person[[0, 14, 15, 16, 17]]
width = 0.0
x0, y0 = candidate[head][:2]
if has_left_eye:
x1, y1 = candidate[left_eye][:2]
d = max(abs(x0 - x1), abs(y0 - y1))
width = max(width, d * 3.0)
if has_right_eye:
x1, y1 = candidate[right_eye][:2]
d = max(abs(x0 - x1), abs(y0 - y1))
width = max(width, d * 3.0)
if has_left_ear:
x1, y1 = candidate[left_ear][:2]
d = max(abs(x0 - x1), abs(y0 - y1))
width = max(width, d * 1.5)
if has_right_ear:
x1, y1 = candidate[right_ear][:2]
d = max(abs(x0 - x1), abs(y0 - y1))
width = max(width, d * 1.5)
x, y = x0, y0
x -= width
y -= width
if x < 0:
x = 0
if y < 0:
y = 0
width1 = width * 2
width2 = width * 2
if x + width > image_width:
width1 = image_width - x
if y + width > image_height:
width2 = image_height - y
width = min(width1, width2)
if width >= 20:
detect_result.append([int(x), int(y), int(width)])
return detect_result
# get max index of 2d array
def npmax(array):
arrayindex = array.argmax(1)
arrayvalue = array.max(1)
i = arrayvalue.argmax()
j = arrayindex[i]
return i, j

80
preprocessing/dwpose/wholebody.py Normal file
View File

@@ -0,0 +1,80 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import cv2
import numpy as np
import onnxruntime as ort
from .onnxdet import inference_detector
from .onnxpose import inference_pose
def HWC3(x):
assert x.dtype == np.uint8
if x.ndim == 2:
x = x[:, :, None]
assert x.ndim == 3
H, W, C = x.shape
assert C == 1 or C == 3 or C == 4
if C == 3:
return x
if C == 1:
return np.concatenate([x, x, x], axis=2)
if C == 4:
color = x[:, :, 0:3].astype(np.float32)
alpha = x[:, :, 3:4].astype(np.float32) / 255.0
y = color * alpha + 255.0 * (1.0 - alpha)
y = y.clip(0, 255).astype(np.uint8)
return y
def resize_image(input_image, resolution):
H, W, C = input_image.shape
H = float(H)
W = float(W)
k = float(resolution) / min(H, W)
H *= k
W *= k
H = int(np.round(H / 64.0)) * 64
W = int(np.round(W / 64.0)) * 64
img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
return img
class Wholebody:
def __init__(self, onnx_det, onnx_pose, device = 'cuda:0'):
providers = ['CPUExecutionProvider'
] if device == 'cpu' else ['CUDAExecutionProvider']
# onnx_det = 'annotator/ckpts/yolox_l.onnx'
# onnx_pose = 'annotator/ckpts/dw-ll_ucoco_384.onnx'
self.session_det = ort.InferenceSession(path_or_bytes=onnx_det, providers=providers)
self.session_pose = ort.InferenceSession(path_or_bytes=onnx_pose, providers=providers)
def __call__(self, ori_img):
det_result = inference_detector(self.session_det, ori_img)
keypoints, scores = inference_pose(self.session_pose, det_result, ori_img)
keypoints_info = np.concatenate(
(keypoints, scores[..., None]), axis=-1)
# compute neck joint
neck = np.mean(keypoints_info[:, [5, 6]], axis=1)
# neck score when visualizing pred
neck[:, 2:4] = np.logical_and(
keypoints_info[:, 5, 2:4] > 0.3,
keypoints_info[:, 6, 2:4] > 0.3).astype(int)
new_keypoints_info = np.insert(
keypoints_info, 17, neck, axis=1)
mmpose_idx = [
17, 6, 8, 10, 7, 9, 12, 14, 16, 13, 15, 2, 1, 4, 3
]
openpose_idx = [
1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17
]
new_keypoints_info[:, openpose_idx] = \
new_keypoints_info[:, mmpose_idx]
keypoints_info = new_keypoints_info
keypoints, scores = keypoints_info[
..., :2], keypoints_info[..., 2]
return keypoints, scores, det_result

35
preprocessing/gray.py Normal file
View File

@@ -0,0 +1,35 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import cv2
import numpy as np
from PIL import Image
import torch
def convert_to_numpy(image):
if isinstance(image, Image.Image):
image = np.array(image)
elif isinstance(image, torch.Tensor):
image = image.detach().cpu().numpy()
elif isinstance(image, np.ndarray):
image = image.copy()
else:
raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
return image
class GrayAnnotator:
def __init__(self, cfg):
pass
def forward(self, image):
image = convert_to_numpy(image)
gray_map = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
return gray_map[..., None].repeat(3, axis=2)
class GrayVideoAnnotator(GrayAnnotator):
def forward(self, frames):
ret_frames = []
for frame in frames:
anno_frame = super().forward(np.array(frame))
ret_frames.append(anno_frame)
return ret_frames

0
preprocessing/matanyone/__init__.py Normal file
View File

698
preprocessing/matanyone/app.py Normal file
View File

@@ -0,0 +1,698 @@
import sys
import os
import json
import time
import psutil
# import ffmpeg
import imageio
from PIL import Image
import cv2
import torch
import numpy as np
import gradio as gr
from .tools.painter import mask_painter
from .tools.interact_tools import SamControler
from .tools.misc import get_device
from .tools.download_util import load_file_from_url
from .utils.get_default_model import get_matanyone_model
from .matanyone.inference.inference_core import InferenceCore
from .matanyone_wrapper import matanyone
arg_device = "cuda"
arg_sam_model_type="vit_h"
arg_mask_save = False
model_loaded = False
model = None
matanyone_model = None
# SAM generator
class MaskGenerator():
def __init__(self, sam_checkpoint, device):
global args_device
args_device = device
self.samcontroler = SamControler(sam_checkpoint, arg_sam_model_type, arg_device)
def first_frame_click(self, image: np.ndarray, points:np.ndarray, labels: np.ndarray, multimask=True):
mask, logit, painted_image = self.samcontroler.first_frame_click(image, points, labels, multimask)
return mask, logit, painted_image
# convert points input to prompt state
def get_prompt(click_state, click_input):
inputs = json.loads(click_input)
points = click_state[0]
labels = click_state[1]
for input in inputs:
points.append(input[:2])
labels.append(input[2])
click_state[0] = points
click_state[1] = labels
prompt = {
"prompt_type":["click"],
"input_point":click_state[0],
"input_label":click_state[1],
"multimask_output":"True",
}
return prompt
def get_frames_from_image(image_input, image_state):
"""
Args:
video_path:str
timestamp:float64
Return
[[0:nearest_frame], [nearest_frame:], nearest_frame]
"""
user_name = time.time()
frames = [image_input] * 2 # hardcode: mimic a video with 2 frames
image_size = (frames[0].shape[0],frames[0].shape[1])
# initialize video_state
image_state = {
"user_name": user_name,
"image_name": "output.png",
"origin_images": frames,
"painted_images": frames.copy(),
"masks": [np.zeros((frames[0].shape[0],frames[0].shape[1]), np.uint8)]*len(frames),
"logits": [None]*len(frames),
"select_frame_number": 0,
"last_frame_numer": 0,
"fps": None
}
image_info = "Image Name: N/A,\nFPS: N/A,\nTotal Frames: {},\nImage Size:{}".format(len(frames), image_size)
model.samcontroler.sam_controler.reset_image()
model.samcontroler.sam_controler.set_image(image_state["origin_images"][0])
return image_state, image_info, image_state["origin_images"][0], \
gr.update(visible=True, maximum=10, value=10), gr.update(visible=True, maximum=len(frames), value=len(frames)), gr.update(visible=False, maximum=len(frames), value=len(frames)), \
gr.update(visible=True), gr.update(visible=True), \
gr.update(visible=True), gr.update(visible=True),\
gr.update(visible=True), gr.update(visible=True), \
gr.update(visible=True), gr.update(visible=False), \
gr.update(visible=False), gr.update(visible=True), \
gr.update(visible=True)
# extract frames from upload video
def get_frames_from_video(video_input, video_state):
"""
Args:
video_path:str
timestamp:float64
Return
[[0:nearest_frame], [nearest_frame:], nearest_frame]
"""
while model == None:
time.sleep(1)
video_path = video_input
frames = []
user_name = time.time()
# extract Audio
# try:
# audio_path = video_input.replace(".mp4", "_audio.wav")
# ffmpeg.input(video_path).output(audio_path, format='wav', acodec='pcm_s16le', ac=2, ar='44100').run(overwrite_output=True, quiet=True)
# except Exception as e:
# print(f"Audio extraction error: {str(e)}")
# audio_path = "" # Set to "" if extraction fails
# print(f'audio_path: {audio_path}')
audio_path = ""
# extract frames
try:
cap = cv2.VideoCapture(video_path)
fps = cap.get(cv2.CAP_PROP_FPS)
while cap.isOpened():
ret, frame = cap.read()
if ret == True:
current_memory_usage = psutil.virtual_memory().percent
frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
if current_memory_usage > 90:
break
else:
break
except (OSError, TypeError, ValueError, KeyError, SyntaxError) as e:
print("read_frame_source:{} error. {}\n".format(video_path, str(e)))
image_size = (frames[0].shape[0],frames[0].shape[1])
# resize if resolution too big
if image_size[0]>=1280 and image_size[0]>=1280:
scale = 1080 / min(image_size)
new_w = int(image_size[1] * scale)
new_h = int(image_size[0] * scale)
# update frames
frames = [cv2.resize(f, (new_w, new_h), interpolation=cv2.INTER_AREA) for f in frames]
# update image_size
image_size = (frames[0].shape[0],frames[0].shape[1])
# initialize video_state
video_state = {
"user_name": user_name,
"video_name": os.path.split(video_path)[-1],
"origin_images": frames,
"painted_images": frames.copy(),
"masks": [np.zeros((frames[0].shape[0],frames[0].shape[1]), np.uint8)]*len(frames),
"logits": [None]*len(frames),
"select_frame_number": 0,
"last_frame_number": 0,
"fps": fps,
"audio": audio_path
}
video_info = "Video Name: {},\nFPS: {},\nTotal Frames: {},\nImage Size:{}".format(video_state["video_name"], round(video_state["fps"], 0), len(frames), image_size)
model.samcontroler.sam_controler.reset_image()
model.samcontroler.sam_controler.set_image(video_state["origin_images"][0])
return video_state, video_info, video_state["origin_images"][0], \
gr.update(visible=True, maximum=len(frames), value=1), gr.update(visible=True, maximum=len(frames), value=len(frames)), gr.update(visible=False, maximum=len(frames), value=len(frames)), \
gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), \
gr.update(visible=True), gr.update(visible=True),\
gr.update(visible=True), gr.update(visible=False), \
gr.update(visible=False), gr.update(visible=False), \
gr.update(visible=False), gr.update(visible=True), \
gr.update(visible=True)
# get the select frame from gradio slider
def select_video_template(image_selection_slider, video_state, interactive_state):
image_selection_slider -= 1
video_state["select_frame_number"] = image_selection_slider
# once select a new template frame, set the image in sam
model.samcontroler.sam_controler.reset_image()
model.samcontroler.sam_controler.set_image(video_state["origin_images"][image_selection_slider])
return video_state["painted_images"][image_selection_slider], video_state, interactive_state
def select_image_template(image_selection_slider, video_state, interactive_state):
image_selection_slider = 0 # fixed for image
video_state["select_frame_number"] = image_selection_slider
# once select a new template frame, set the image in sam
model.samcontroler.sam_controler.reset_image()
model.samcontroler.sam_controler.set_image(video_state["origin_images"][image_selection_slider])
return video_state["painted_images"][image_selection_slider], video_state, interactive_state
# set the tracking end frame
def get_end_number(track_pause_number_slider, video_state, interactive_state):
interactive_state["track_end_number"] = track_pause_number_slider
return video_state["painted_images"][track_pause_number_slider],interactive_state
# use sam to get the mask
def sam_refine(video_state, point_prompt, click_state, interactive_state, evt:gr.SelectData ): #
"""
Args:
template_frame: PIL.Image
point_prompt: flag for positive or negative button click
click_state: [[points], [labels]]
"""
if point_prompt == "Positive":
coordinate = "[[{},{},1]]".format(evt.index[0], evt.index[1])
interactive_state["positive_click_times"] += 1
else:
coordinate = "[[{},{},0]]".format(evt.index[0], evt.index[1])
interactive_state["negative_click_times"] += 1
# prompt for sam model
model.samcontroler.sam_controler.reset_image()
model.samcontroler.sam_controler.set_image(video_state["origin_images"][video_state["select_frame_number"]])
prompt = get_prompt(click_state=click_state, click_input=coordinate)
mask, logit, painted_image = model.first_frame_click(
image=video_state["origin_images"][video_state["select_frame_number"]],
points=np.array(prompt["input_point"]),
labels=np.array(prompt["input_label"]),
multimask=prompt["multimask_output"],
)
video_state["masks"][video_state["select_frame_number"]] = mask
video_state["logits"][video_state["select_frame_number"]] = logit
video_state["painted_images"][video_state["select_frame_number"]] = painted_image
return painted_image, video_state, interactive_state
def add_multi_mask(video_state, interactive_state, mask_dropdown):
mask = video_state["masks"][video_state["select_frame_number"]]
interactive_state["multi_mask"]["masks"].append(mask)
interactive_state["multi_mask"]["mask_names"].append("mask_{:03d}".format(len(interactive_state["multi_mask"]["masks"])))
mask_dropdown.append("mask_{:03d}".format(len(interactive_state["multi_mask"]["masks"])))
select_frame = show_mask(video_state, interactive_state, mask_dropdown)
return interactive_state, gr.update(choices=interactive_state["multi_mask"]["mask_names"], value=mask_dropdown), select_frame, [[],[]]
def clear_click(video_state, click_state):
click_state = [[],[]]
template_frame = video_state["origin_images"][video_state["select_frame_number"]]
return template_frame, click_state
def remove_multi_mask(interactive_state, mask_dropdown):
interactive_state["multi_mask"]["mask_names"]= []
interactive_state["multi_mask"]["masks"] = []
return interactive_state, gr.update(choices=[],value=[])
def show_mask(video_state, interactive_state, mask_dropdown):
mask_dropdown.sort()
if video_state["origin_images"]:
select_frame = video_state["origin_images"][video_state["select_frame_number"]]
for i in range(len(mask_dropdown)):
mask_number = int(mask_dropdown[i].split("_")[1]) - 1
mask = interactive_state["multi_mask"]["masks"][mask_number]
select_frame = mask_painter(select_frame, mask.astype('uint8'), mask_color=mask_number+2)
return select_frame
def save_video(frames, output_path, fps):
writer = imageio.get_writer( output_path, fps=fps, codec='libx264', quality=8)
for frame in frames:
writer.append_data(frame)
writer.close()
return output_path
# video matting
def video_matting(video_state, end_slider, matting_type, interactive_state, mask_dropdown, erode_kernel_size, dilate_kernel_size):
matanyone_processor = InferenceCore(matanyone_model, cfg=matanyone_model.cfg)
# if interactive_state["track_end_number"]:
# following_frames = video_state["origin_images"][video_state["select_frame_number"]:interactive_state["track_end_number"]]
# else:
end_slider = max(video_state["select_frame_number"] +1, end_slider)
following_frames = video_state["origin_images"][video_state["select_frame_number"]: end_slider]
if interactive_state["multi_mask"]["masks"]:
if len(mask_dropdown) == 0:
mask_dropdown = ["mask_001"]
mask_dropdown.sort()
template_mask = interactive_state["multi_mask"]["masks"][int(mask_dropdown[0].split("_")[1]) - 1] * (int(mask_dropdown[0].split("_")[1]))
for i in range(1,len(mask_dropdown)):
mask_number = int(mask_dropdown[i].split("_")[1]) - 1
template_mask = np.clip(template_mask+interactive_state["multi_mask"]["masks"][mask_number]*(mask_number+1), 0, mask_number+1)
video_state["masks"][video_state["select_frame_number"]]= template_mask
else:
template_mask = video_state["masks"][video_state["select_frame_number"]]
fps = video_state["fps"]
audio_path = video_state["audio"]
# operation error
if len(np.unique(template_mask))==1:
template_mask[0][0]=1
foreground, alpha = matanyone(matanyone_processor, following_frames, template_mask*255, r_erode=erode_kernel_size, r_dilate=dilate_kernel_size)
output_frames = []
foreground_mat = matting_type == "Foreground"
for frame_origin, frame_alpha in zip(following_frames, alpha):
if foreground_mat:
frame_alpha[frame_alpha > 127] = 255
frame_alpha[frame_alpha <= 127] = 0
else:
frame_temp = frame_alpha.copy()
frame_alpha[frame_temp > 127] = 0
frame_alpha[frame_temp <= 127] = 255
output_frame = np.bitwise_and(frame_origin, 255-frame_alpha)
frame_grey = frame_alpha.copy()
frame_grey[frame_alpha == 255] = 127
output_frame += frame_grey
output_frames.append(output_frame)
foreground = output_frames
if not os.path.exists("mask_outputs"):
os.makedirs("mask_outputs")
file_name= video_state["video_name"]
file_name = ".".join(file_name.split(".")[:-1])
foreground_output = save_video(foreground, output_path="./mask_outputs/{}_fg.mp4".format(file_name), fps=fps)
# foreground_output = generate_video_from_frames(foreground, output_path="./results/{}_fg.mp4".format(video_state["video_name"]), fps=fps, audio_path=audio_path) # import video_input to name the output video
alpha_output = save_video(alpha, output_path="./mask_outputs/{}_alpha.mp4".format(file_name), fps=fps)
# alpha_output = generate_video_from_frames(alpha, output_path="./results/{}_alpha.mp4".format(video_state["video_name"]), fps=fps, gray2rgb=True, audio_path=audio_path) # import video_input to name the output video
return foreground_output, alpha_output, gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
def show_outputs():
return gr.update(visible=True), gr.update(visible=True)
def add_audio_to_video(video_path, audio_path, output_path):
try:
video_input = ffmpeg.input(video_path)
audio_input = ffmpeg.input(audio_path)
_ = (
ffmpeg
.output(video_input, audio_input, output_path, vcodec="copy", acodec="aac")
.run(overwrite_output=True, capture_stdout=True, capture_stderr=True)
)
return output_path
except ffmpeg.Error as e:
print(f"FFmpeg error:\n{e.stderr.decode()}")
return None
def generate_video_from_frames(frames, output_path, fps=30, gray2rgb=False, audio_path=""):
"""
Generates a video from a list of frames.
Args:
frames (list of numpy arrays): The frames to include in the video.
output_path (str): The path to save the generated video.
fps (int, optional): The frame rate of the output video. Defaults to 30.
"""
frames = torch.from_numpy(np.asarray(frames))
_, h, w, _ = frames.shape
if gray2rgb:
frames = np.repeat(frames, 3, axis=3)
if not os.path.exists(os.path.dirname(output_path)):
os.makedirs(os.path.dirname(output_path))
video_temp_path = output_path.replace(".mp4", "_temp.mp4")
# resize back to ensure input resolution
imageio.mimwrite(video_temp_path, frames, fps=fps, quality=7,
codec='libx264', ffmpeg_params=["-vf", f"scale={w}:{h}"])
# add audio to video if audio path exists
if audio_path != "" and os.path.exists(audio_path):
output_path = add_audio_to_video(video_temp_path, audio_path, output_path)
os.remove(video_temp_path)
return output_path
else:
return video_temp_path
# reset all states for a new input
def restart():
return {
"user_name": "",
"video_name": "",
"origin_images": None,
"painted_images": None,
"masks": None,
"inpaint_masks": None,
"logits": None,
"select_frame_number": 0,
"fps": 30
}, {
"inference_times": 0,
"negative_click_times" : 0,
"positive_click_times": 0,
"mask_save": arg_mask_save,
"multi_mask": {
"mask_names": [],
"masks": []
},
"track_end_number": None,
}, [[],[]], None, None, \
gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False),\
gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), \
gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), \
gr.update(visible=False), gr.update(visible=False, choices=[], value=[]), "", gr.update(visible=False)
def load_unload_models(selected):
global model_loaded
global model
global matanyone_model
if selected:
if model_loaded:
model.samcontroler.sam_controler.model.to(arg_device)
matanyone_model.to(arg_device)
else:
# args, defined in track_anything.py
sam_checkpoint_url_dict = {
'vit_h': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
'vit_l': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth",
'vit_b': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
}
# os.path.join('.')
from mmgp import offload
# sam_checkpoint = load_file_from_url(sam_checkpoint_url_dict[arg_sam_model_type], ".")
sam_checkpoint = None
transfer_stream = torch.cuda.Stream()
with torch.cuda.stream(transfer_stream):
# initialize sams
model = MaskGenerator(sam_checkpoint, arg_device)
from .matanyone.model.matanyone import MatAnyone
matanyone_model = MatAnyone.from_pretrained("PeiqingYang/MatAnyone")
# pipe ={"mat" : matanyone_model, "sam" :model.samcontroler.sam_controler.model }
# offload.profile(pipe)
matanyone_model = matanyone_model.to(arg_device).eval()
matanyone_processor = InferenceCore(matanyone_model, cfg=matanyone_model.cfg)
model_loaded = True
else:
import gc
model.samcontroler.sam_controler.model.to("cpu")
matanyone_model.to("cpu")
gc.collect()
torch.cuda.empty_cache()
def get_vmc_event_handler():
return load_unload_models
def export_to_vace_video_input(foreground_video_output):
gr.Info("Masked Video Input transferred to Vace For Inpainting")
return "V#" + str(time.time()), foreground_video_output
def export_to_vace_video_mask(foreground_video_output, alpha_video_output):
gr.Info("Masked Video Input and Full Mask transferred to Vace For Inpainting")
return "MV#" + str(time.time()), foreground_video_output, alpha_video_output
def teleport_to_vace():
return gr.Tabs(selected="video_gen"), gr.Dropdown(value="vace_1.3B")
def display(tabs, model_choice, vace_video_input, vace_video_mask, video_prompt_video_guide_trigger):
# my_tab.select(fn=load_unload_models, inputs=[], outputs=[])
media_url = "https://github.com/pq-yang/MatAnyone/releases/download/media/"
# download assets
gr.Markdown("Mast Edition is provided by MatAnyone")
with gr.Column( visible=True):
with gr.Row():
with gr.Accordion("Video Tutorial (click to expand)", open=False, elem_classes="custom-bg"):
with gr.Row():
with gr.Column():
gr.Markdown("### Case 1: Single Target")
gr.Video(value="preprocessing/matanyone/tutorial_single_target.mp4", elem_classes="video")
with gr.Column():
gr.Markdown("### Case 2: Multiple Targets")
gr.Video(value="preprocessing/matanyone/tutorial_multi_targets.mp4", elem_classes="video")
click_state = gr.State([[],[]])
interactive_state = gr.State({
"inference_times": 0,
"negative_click_times" : 0,
"positive_click_times": 0,
"mask_save": arg_mask_save,
"multi_mask": {
"mask_names": [],
"masks": []
},
"track_end_number": None,
}
)
video_state = gr.State(
{
"user_name": "",
"video_name": "",
"origin_images": None,
"painted_images": None,
"masks": None,
"inpaint_masks": None,
"logits": None,
"select_frame_number": 0,
"fps": 16,
"audio": "",
}
)
with gr.Column( visible=True):
with gr.Row():
with gr.Accordion('MatAnyone Settings (click to expand)', open=False):
with gr.Row():
erode_kernel_size = gr.Slider(label='Erode Kernel Size',
minimum=0,
maximum=30,
step=1,
value=10,
info="Erosion on the added mask",
interactive=True)
dilate_kernel_size = gr.Slider(label='Dilate Kernel Size',
minimum=0,
maximum=30,
step=1,
value=10,
info="Dilation on the added mask",
interactive=True)
with gr.Row():
image_selection_slider = gr.Slider(minimum=1, maximum=100, step=1, value=1, label="Start Frame", info="Choose the start frame for target assignment and video matting", visible=False)
end_selection_slider = gr.Slider(minimum=1, maximum=300, step=1, value=81, label="Last Frame to Process", info="Last Frame to Process", visible=False)
track_pause_number_slider = gr.Slider(minimum=1, maximum=100, step=1, value=1, label="End frame", visible=False)
with gr.Row():
point_prompt = gr.Radio(
choices=["Positive", "Negative"],
value="Positive",
label="Point Prompt",
info="Click to add positive or negative point for target mask",
interactive=True,
visible=False,
min_width=100,
scale=1)
matting_type = gr.Radio(
choices=["Foreground", "Background"],
value="Foreground",
label="Matting Type",
info="Type of Video Matting to Generate",
interactive=True,
visible=False,
min_width=100,
scale=1)
mask_dropdown = gr.Dropdown(multiselect=True, value=[], label="Mask Selection", info="Choose 1~all mask(s) added in Step 2", visible=False, scale=2)
gr.Markdown("---")
with gr.Column():
# input video
with gr.Row(equal_height=True):
with gr.Column(scale=2):
gr.Markdown("## Step1: Upload video")
with gr.Column(scale=2):
step2_title = gr.Markdown("## Step2: Add masks <small>(Several clicks then **`Add Mask`** <u>one by one</u>)</small>", visible=False)
with gr.Row(equal_height=True):
with gr.Column(scale=2):
video_input = gr.Video(label="Input Video", elem_classes="video")
extract_frames_button = gr.Button(value="Load Video", interactive=True, elem_classes="new_button")
with gr.Column(scale=2):
video_info = gr.Textbox(label="Video Info", visible=False)
template_frame = gr.Image(label="Start Frame", type="pil",interactive=True, elem_id="template_frame", visible=False, elem_classes="image")
with gr.Row():
clear_button_click = gr.Button(value="Clear Clicks", interactive=True, visible=False, min_width=100)
add_mask_button = gr.Button(value="Set Mask", interactive=True, visible=False, min_width=100)
remove_mask_button = gr.Button(value="Remove Mask", interactive=True, visible=False, min_width=100) # no use
matting_button = gr.Button(value="Generate Video Matting", interactive=True, visible=False, min_width=100)
with gr.Row():
gr.Markdown("")
# output video
with gr.Column() as output_row: #equal_height=True
with gr.Row():
with gr.Column(scale=2):
foreground_video_output = gr.Video(label="Masked Video Output", visible=False, elem_classes="video")
foreground_output_button = gr.Button(value="Black & White Video Output", visible=False, elem_classes="new_button")
with gr.Column(scale=2):
alpha_video_output = gr.Video(label="B & W Mask Video Output", visible=False, elem_classes="video")
alpha_output_button = gr.Button(value="Alpha Mask Output", visible=False, elem_classes="new_button")
with gr.Row():
with gr.Row(visible= False):
export_to_vace_video_input_btn = gr.Button("Export to Vace Video Input Video For Inpainting", visible= False)
with gr.Row(visible= True):
export_to_vace_video_mask_btn = gr.Button("Export to Vace Video Input and Video Mask", visible= False)
export_to_vace_video_input_btn.click(fn=export_to_vace_video_input, inputs= [foreground_video_output], outputs= [video_prompt_video_guide_trigger, vace_video_input])
export_to_vace_video_mask_btn.click(fn=export_to_vace_video_mask, inputs= [foreground_video_output, alpha_video_output], outputs= [video_prompt_video_guide_trigger, vace_video_input, vace_video_mask]).then(
fn=teleport_to_vace, inputs=[], outputs=[tabs, model_choice])
# first step: get the video information
extract_frames_button.click(
fn=get_frames_from_video,
inputs=[
video_input, video_state
],
outputs=[video_state, video_info, template_frame,
image_selection_slider, end_selection_slider, track_pause_number_slider, point_prompt, matting_type, clear_button_click, add_mask_button, matting_button, template_frame,
foreground_video_output, alpha_video_output, foreground_output_button, alpha_output_button, mask_dropdown, step2_title]
)
# second step: select images from slider
image_selection_slider.release(fn=select_video_template,
inputs=[image_selection_slider, video_state, interactive_state],
outputs=[template_frame, video_state, interactive_state], api_name="select_image")
track_pause_number_slider.release(fn=get_end_number,
inputs=[track_pause_number_slider, video_state, interactive_state],
outputs=[template_frame, interactive_state], api_name="end_image")
# click select image to get mask using sam
template_frame.select(
fn=sam_refine,
inputs=[video_state, point_prompt, click_state, interactive_state],
outputs=[template_frame, video_state, interactive_state]
)
# add different mask
add_mask_button.click(
fn=add_multi_mask,
inputs=[video_state, interactive_state, mask_dropdown],
outputs=[interactive_state, mask_dropdown, template_frame, click_state]
)
remove_mask_button.click(
fn=remove_multi_mask,
inputs=[interactive_state, mask_dropdown],
outputs=[interactive_state, mask_dropdown]
)
# video matting
matting_button.click(
fn=show_outputs,
inputs=[],
outputs=[foreground_video_output, alpha_video_output]).then(
fn=video_matting,
inputs=[video_state, end_selection_slider, matting_type, interactive_state, mask_dropdown, erode_kernel_size, dilate_kernel_size],
outputs=[foreground_video_output, alpha_video_output,foreground_video_output, alpha_video_output, export_to_vace_video_input_btn, export_to_vace_video_mask_btn]
)
# click to get mask
mask_dropdown.change(
fn=show_mask,
inputs=[video_state, interactive_state, mask_dropdown],
outputs=[template_frame]
)
# clear input
video_input.change(
fn=restart,
inputs=[],
outputs=[
video_state,
interactive_state,
click_state,
foreground_video_output, alpha_video_output,
template_frame,
image_selection_slider, end_selection_slider, track_pause_number_slider,point_prompt, export_to_vace_video_input_btn, export_to_vace_video_mask_btn, matting_type, clear_button_click,
add_mask_button, matting_button, template_frame, foreground_video_output, alpha_video_output, remove_mask_button, foreground_output_button, alpha_output_button, mask_dropdown, video_info, step2_title
],
queue=False,
show_progress=False)
video_input.clear(
fn=restart,
inputs=[],
outputs=[
video_state,
interactive_state,
click_state,
foreground_video_output, alpha_video_output,
template_frame,
image_selection_slider , end_selection_slider, track_pause_number_slider,point_prompt, export_to_vace_video_input_btn, export_to_vace_video_mask_btn, matting_type, clear_button_click,
add_mask_button, matting_button, template_frame, foreground_video_output, alpha_video_output, remove_mask_button, foreground_output_button, alpha_output_button, mask_dropdown, video_info, step2_title
],
queue=False,
show_progress=False)
# points clear
clear_button_click.click(
fn = clear_click,
inputs = [video_state, click_state,],
outputs = [template_frame,click_state],
)

0
preprocessing/matanyone/matanyone/config/__init__.py Normal file
View File

47
preprocessing/matanyone/matanyone/config/eval_matanyone_config.yaml Normal file
View File

@@ -0,0 +1,47 @@
defaults:
- _self_
- model: base
- override hydra/job_logging: custom-no-rank.yaml
hydra:
run:
dir: ../output/${exp_id}/${dataset}
output_subdir: ${now:%Y-%m-%d_%H-%M-%S}-hydra
amp: False
weights: pretrained_models/matanyone.pth # default (can be modified from outside)
output_dir: null # defaults to run_dir; specify this to override
flip_aug: False
# maximum shortest side of the input; -1 means no resizing
# With eval_vos.py, we usually just use the dataset's size (resizing done in dataloader)
# this parameter is added for the sole purpose for the GUI in the current codebase
# InferenceCore will downsize the input and restore the output to the original size if needed
# if you are using this code for some other project, you can also utilize this parameter
max_internal_size: -1
# these parameters, when set, override the dataset's default; useful for debugging
save_all: True
use_all_masks: False
use_long_term: False
mem_every: 5
# only relevant when long_term is not enabled
max_mem_frames: 5
# only relevant when long_term is enabled
long_term:
count_usage: True
max_mem_frames: 10
min_mem_frames: 5
num_prototypes: 128
max_num_tokens: 10000
buffer_tokens: 2000
top_k: 30
stagger_updates: 5
chunk_size: -1 # number of objects to process in parallel; -1 means unlimited
save_scores: False
save_aux: False
visualize: False

22
preprocessing/matanyone/matanyone/config/hydra/job_logging/custom-no-rank.yaml Normal file
View File

@@ -0,0 +1,22 @@
# python logging configuration for tasks
version: 1
formatters:
simple:
format: '[%(asctime)s][%(levelname)s] - %(message)s'
datefmt: '%Y-%m-%d %H:%M:%S'
handlers:
console:
class: logging.StreamHandler
formatter: simple
stream: ext://sys.stdout
file:
class: logging.FileHandler
formatter: simple
# absolute file path
filename: ${hydra.runtime.output_dir}/${now:%Y-%m-%d_%H-%M-%S}-eval.log
mode: w
root:
level: INFO
handlers: [console, file]
disable_existing_loggers: false

22
preprocessing/matanyone/matanyone/config/hydra/job_logging/custom.yaml Normal file
View File

@@ -0,0 +1,22 @@
# python logging configuration for tasks
version: 1
formatters:
simple:
format: '[%(asctime)s][%(levelname)s][r${oc.env:LOCAL_RANK}] - %(message)s'
datefmt: '%Y-%m-%d %H:%M:%S'
handlers:
console:
class: logging.StreamHandler
formatter: simple
stream: ext://sys.stdout
file:
class: logging.FileHandler
formatter: simple
# absolute file path
filename: ${hydra.runtime.output_dir}/${now:%Y-%m-%d_%H-%M-%S}-rank${oc.env:LOCAL_RANK}.log
mode: w
root:
level: INFO
handlers: [console, file]
disable_existing_loggers: false

58
preprocessing/matanyone/matanyone/config/model/base.yaml Normal file
View File

@@ -0,0 +1,58 @@
pixel_mean: [0.485, 0.456, 0.406]
pixel_std: [0.229, 0.224, 0.225]
pixel_dim: 256
key_dim: 64
value_dim: 256
sensory_dim: 256
embed_dim: 256
pixel_encoder:
type: resnet50
ms_dims: [1024, 512, 256, 64, 3] # f16, f8, f4, f2, f1
mask_encoder:
type: resnet18
final_dim: 256
pixel_pe_scale: 32
pixel_pe_temperature: 128
object_transformer:
embed_dim: ${model.embed_dim}
ff_dim: 2048
num_heads: 8
num_blocks: 3
num_queries: 16
read_from_pixel:
input_norm: False
input_add_pe: False
add_pe_to_qkv: [True, True, False]
read_from_past:
add_pe_to_qkv: [True, True, False]
read_from_memory:
add_pe_to_qkv: [True, True, False]
read_from_query:
add_pe_to_qkv: [True, True, False]
output_norm: False
query_self_attention:
add_pe_to_qkv: [True, True, False]
pixel_self_attention:
add_pe_to_qkv: [True, True, False]
object_summarizer:
embed_dim: ${model.object_transformer.embed_dim}
num_summaries: ${model.object_transformer.num_queries}
add_pe: True
aux_loss:
sensory:
enabled: True
weight: 0.01
query:
enabled: True
weight: 0.01
mask_decoder:
# first value must equal embed_dim
up_dims: [256, 128, 128, 64, 16]

0
preprocessing/matanyone/matanyone/inference/__init__.py Normal file
View File

56
preprocessing/matanyone/matanyone/inference/image_feature_store.py Normal file
View File

@@ -0,0 +1,56 @@
import warnings
from typing import Iterable
import torch
from ..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')

406
preprocessing/matanyone/matanyone/inference/inference_core.py Normal file
View File

@@ -0,0 +1,406 @@
from typing import List, Optional, Iterable
import logging
from omegaconf import DictConfig
import numpy as np
import torch
import torch.nn.functional as F
from .memory_manager import MemoryManager
from .object_manager import ObjectManager
from .image_feature_store import ImageFeatureStore
from ..model.matanyone import MatAnyone
from ...utils.tensor_utils import pad_divide_by, unpad, aggregate
log = logging.getLogger()
class InferenceCore:
def __init__(self,
network: MatAnyone,
cfg: DictConfig,
*,
image_feature_store: ImageFeatureStore = None):
self.network = network
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

348
preprocessing/matanyone/matanyone/inference/kv_memory_store.py Normal file
View File

@@ -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

453
preprocessing/matanyone/matanyone/inference/memory_manager.py Normal file
View File

@@ -0,0 +1,453 @@
import logging
from omegaconf import DictConfig
from typing import List, Dict
import torch
from .object_manager import ObjectManager
from .kv_memory_store import KeyValueMemoryStore
from ..model.matanyone import MatAnyone
from ..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 = {}

24
preprocessing/matanyone/matanyone/inference/object_info.py Normal file
View File

@@ -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})'

149
preprocessing/matanyone/matanyone/inference/object_manager.py Normal file
View File

@@ -0,0 +1,149 @@
from typing import Union, List, Dict
import torch
from .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]]

0
preprocessing/matanyone/matanyone/inference/utils/__init__.py Normal file
View File

30
preprocessing/matanyone/matanyone/inference/utils/args_utils.py Normal file
View File

@@ -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

0
preprocessing/matanyone/matanyone/model/__init__.py Normal file
View File

93
preprocessing/matanyone/matanyone/model/aux_modules.py Normal file
View File

@@ -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 .group_modules import GConv2d
from ...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

365
preprocessing/matanyone/matanyone/model/big_modules.py Normal file
View File

@@ -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 .group_modules import MainToGroupDistributor, GroupFeatureFusionBlock, GConv2d
from .utils import resnet
from .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"):
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

39
preprocessing/matanyone/matanyone/model/channel_attn.py Normal file
View File

@@ -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

126
preprocessing/matanyone/matanyone/model/group_modules.py Normal file
View File

@@ -0,0 +1,126 @@
from typing import Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
from .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

333
preprocessing/matanyone/matanyone/model/matanyone.py Normal file
View File

@@ -0,0 +1,333 @@
from typing import List, Dict, Iterable
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 .big_modules import PixelEncoder, UncertPred, KeyProjection, MaskEncoder, PixelFeatureFuser, MaskDecoder
from .aux_modules import AuxComputer
from .utils.memory_utils import get_affinity, readout
from .transformer.object_transformer import QueryTransformer
from .transformer.object_summarizer import ObjectSummarizer
from ...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) -> (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) -> (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) -> (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.cuda.amp.autocast(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) -> (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) -> (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) -> (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

149
preprocessing/matanyone/matanyone/model/modules.py Normal file
View File

@@ -0,0 +1,149 @@
from typing import List, Iterable
import torch
import torch.nn as nn
import torch.nn.functional as F
from .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"):
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"):
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"):
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

0
preprocessing/matanyone/matanyone/model/transformer/__init__.py Normal file
View File

89
preprocessing/matanyone/matanyone/model/transformer/object_summarizer.py Normal file
View File

@@ -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 .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"):
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

206
preprocessing/matanyone/matanyone/model/transformer/object_transformer.py Normal file
View File

@@ -0,0 +1,206 @@
from typing import Dict, Optional
from omegaconf import DictConfig
import torch
import torch.nn as nn
from ..group_modules import GConv2d
from ....utils.tensor_utils import aggregate
from .positional_encoding import PositionalEncoding
from .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

108
preprocessing/matanyone/matanyone/model/transformer/positional_encoding.py Normal file
View File

@@ -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, :])

161
preprocessing/matanyone/matanyone/model/transformer/transformer_layers.py Normal file
View File

@@ -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 ...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))

0
preprocessing/matanyone/matanyone/model/utils/__init__.py Normal file
View File

107
preprocessing/matanyone/matanyone/model/utils/memory_utils.py Normal file
View File

@@ -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

72
preprocessing/matanyone/matanyone/model/utils/parameter_groups.py Normal file
View File

@@ -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

179
preprocessing/matanyone/matanyone/model/utils/resnet.py Normal file
View File

@@ -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

73
preprocessing/matanyone/matanyone_wrapper.py Normal file
View File

@@ -0,0 +1,73 @@
import tqdm
import torch
from torchvision.transforms.functional import to_tensor
import numpy as np
import random
import cv2
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)
@torch.inference_mode()
@torch.amp.autocast('cuda')
def matanyone(processor, frames_np, mask, r_erode=0, r_dilate=0, n_warmup=10):
"""
Args:
frames_np: [(H,W,C)]*n, uint8
mask: (H,W), uint8
Outputs:
com: [(H,W,C)]*n, uint8
pha: [(H,W,C)]*n, uint8
"""
# print(f'===== [r_erode] {r_erode}; [r_dilate] {r_dilate} =====')
bgr = (np.array([120, 255, 155], dtype=np.float32)/255).reshape((1, 1, 3))
objects = [1]
# [optional] erode & dilate on given seg mask
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()
frames_np = [frames_np[0]]* n_warmup + frames_np
frames = []
phas = []
for ti, frame_single in tqdm.tqdm(enumerate(frames_np)):
image = to_tensor(frame_single).cuda().float()
if ti == 0:
output_prob = processor.step(image, mask, objects=objects) # encode given mask
output_prob = processor.step(image, first_frame_pred=True) # clear past memory for warmup frames
else:
if ti <= n_warmup:
output_prob = processor.step(image, first_frame_pred=True) # clear past memory for warmup frames
else:
output_prob = processor.step(image)
# convert output probabilities to an object mask
mask = processor.output_prob_to_mask(output_prob)
pha = mask.unsqueeze(2).cpu().numpy()
com_np = frame_single / 255. * pha + bgr * (1 - pha)
# DONOT save the warmup frames
if ti > (n_warmup-1):
frames.append((com_np*255).astype(np.uint8))
phas.append((pha*255).astype(np.uint8))
return frames, phas

0
preprocessing/matanyone/tools/__init__.py Normal file
View File

141
preprocessing/matanyone/tools/base_segmenter.py Normal file
View File

@@ -0,0 +1,141 @@
import time
import torch
import cv2
from PIL import Image, ImageDraw, ImageOps
import numpy as np
from typing import Union
from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator
import matplotlib.pyplot as plt
import PIL
from .mask_painter import mask_painter
class BaseSegmenter:
def __init__(self, SAM_checkpoint, model_type, device='cuda:0'):
"""
device: model device
SAM_checkpoint: path of SAM checkpoint
model_type: vit_b, vit_l, vit_h
"""
print(f"Initializing BaseSegmenter to {device}")
assert model_type in ['vit_b', 'vit_l', 'vit_h'], 'model_type must be vit_b, vit_l, or vit_h'
self.device = device
# SAM_checkpoint = None
self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32
from accelerate import init_empty_weights
# self.model = sam_model_registry[model_type](checkpoint=SAM_checkpoint)
with init_empty_weights():
self.model = sam_model_registry[model_type](checkpoint=SAM_checkpoint)
from mmgp import offload
# self.model.to(torch.float16)
# offload.save_model(self.model, "ckpts/mask/sam_vit_h_4b8939_fp16.safetensors")
offload.load_model_data(self.model, "ckpts/mask/sam_vit_h_4b8939_fp16.safetensors")
self.model.to(torch.float32) # need to be optimized, if not f32 crappy precision
self.model.to(device=self.device)
self.predictor = SamPredictor(self.model)
self.embedded = False
@torch.no_grad()
def set_image(self, image: np.ndarray):
# PIL.open(image_path) 3channel: RGB
# image embedding: avoid encode the same image multiple times
self.orignal_image = image
if self.embedded:
print('repeat embedding, please reset_image.')
return
self.predictor.set_image(image)
self.embedded = True
return
@torch.no_grad()
def reset_image(self):
# reset image embeding
self.predictor.reset_image()
self.embedded = False
def predict(self, prompts, mode, multimask=True):
"""
image: numpy array, h, w, 3
prompts: dictionary, 3 keys: 'point_coords', 'point_labels', 'mask_input'
prompts['point_coords']: numpy array [N,2]
prompts['point_labels']: numpy array [1,N]
prompts['mask_input']: numpy array [1,256,256]
mode: 'point' (points only), 'mask' (mask only), 'both' (consider both)
mask_outputs: True (return 3 masks), False (return 1 mask only)
whem mask_outputs=True, mask_input=logits[np.argmax(scores), :, :][None, :, :]
"""
assert self.embedded, 'prediction is called before set_image (feature embedding).'
assert mode in ['point', 'mask', 'both'], 'mode must be point, mask, or both'
with torch.autocast(device_type='cuda', dtype=torch.float16):
if mode == 'point':
masks, scores, logits = self.predictor.predict(point_coords=prompts['point_coords'],
point_labels=prompts['point_labels'],
multimask_output=multimask)
elif mode == 'mask':
masks, scores, logits = self.predictor.predict(mask_input=prompts['mask_input'],
multimask_output=multimask)
elif mode == 'both': # both
masks, scores, logits = self.predictor.predict(point_coords=prompts['point_coords'],
point_labels=prompts['point_labels'],
mask_input=prompts['mask_input'],
multimask_output=multimask)
else:
raise("Not implement now!")
# masks (n, h, w), scores (n,), logits (n, 256, 256)
return masks, scores, logits
if __name__ == "__main__":
# load and show an image
image = cv2.imread('/hhd3/gaoshang/truck.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # numpy array (h, w, 3)
# initialise BaseSegmenter
SAM_checkpoint= '/ssd1/gaomingqi/checkpoints/sam_vit_h_4b8939.pth'
model_type = 'vit_h'
device = "cuda:4"
base_segmenter = BaseSegmenter(SAM_checkpoint=SAM_checkpoint, model_type=model_type, device=device)
# image embedding (once embedded, multiple prompts can be applied)
base_segmenter.set_image(image)
# examples
# point only ------------------------
mode = 'point'
prompts = {
'point_coords': np.array([[500, 375], [1125, 625]]),
'point_labels': np.array([1, 1]),
}
masks, scores, logits = base_segmenter.predict(prompts, mode, multimask=False) # masks (n, h, w), scores (n,), logits (n, 256, 256)
painted_image = mask_painter(image, masks[np.argmax(scores)].astype('uint8'), background_alpha=0.8)
painted_image = cv2.cvtColor(painted_image, cv2.COLOR_RGB2BGR) # numpy array (h, w, 3)
cv2.imwrite('/hhd3/gaoshang/truck_point.jpg', painted_image)
# both ------------------------
mode = 'both'
mask_input = logits[np.argmax(scores), :, :]
prompts = {'mask_input': mask_input [None, :, :]}
prompts = {
'point_coords': np.array([[500, 375], [1125, 625]]),
'point_labels': np.array([1, 0]),
'mask_input': mask_input[None, :, :]
}
masks, scores, logits = base_segmenter.predict(prompts, mode, multimask=True) # masks (n, h, w), scores (n,), logits (n, 256, 256)
painted_image = mask_painter(image, masks[np.argmax(scores)].astype('uint8'), background_alpha=0.8)
painted_image = cv2.cvtColor(painted_image, cv2.COLOR_RGB2BGR) # numpy array (h, w, 3)
cv2.imwrite('/hhd3/gaoshang/truck_both.jpg', painted_image)
# mask only ------------------------
mode = 'mask'
mask_input = logits[np.argmax(scores), :, :]
prompts = {'mask_input': mask_input[None, :, :]}
masks, scores, logits = base_segmenter.predict(prompts, mode, multimask=True) # masks (n, h, w), scores (n,), logits (n, 256, 256)
painted_image = mask_painter(image, masks[np.argmax(scores)].astype('uint8'), background_alpha=0.8)
painted_image = cv2.cvtColor(painted_image, cv2.COLOR_RGB2BGR) # numpy array (h, w, 3)
cv2.imwrite('/hhd3/gaoshang/truck_mask.jpg', painted_image)

109
preprocessing/matanyone/tools/download_util.py Normal file
View File

@@ -0,0 +1,109 @@
import math
import os
import requests
from torch.hub import download_url_to_file, get_dir
from tqdm import tqdm
from urllib.parse import urlparse
def sizeof_fmt(size, suffix='B'):
"""Get human readable file size.
Args:
size (int): File size.
suffix (str): Suffix. Default: 'B'.
Return:
str: Formated file siz.
"""
for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:
if abs(size) < 1024.0:
return f'{size:3.1f} {unit}{suffix}'
size /= 1024.0
return f'{size:3.1f} Y{suffix}'
def download_file_from_google_drive(file_id, save_path):
"""Download files from google drive.
Ref:
https://stackoverflow.com/questions/25010369/wget-curl-large-file-from-google-drive # noqa E501
Args:
file_id (str): File id.
save_path (str): Save path.
"""
session = requests.Session()
URL = 'https://docs.google.com/uc?export=download'
params = {'id': file_id}
response = session.get(URL, params=params, stream=True)
token = get_confirm_token(response)
if token:
params['confirm'] = token
response = session.get(URL, params=params, stream=True)
# get file size
response_file_size = session.get(URL, params=params, stream=True, headers={'Range': 'bytes=0-2'})
print(response_file_size)
if 'Content-Range' in response_file_size.headers:
file_size = int(response_file_size.headers['Content-Range'].split('/')[1])
else:
file_size = None
save_response_content(response, save_path, file_size)
def get_confirm_token(response):
for key, value in response.cookies.items():
if key.startswith('download_warning'):
return value
return None
def save_response_content(response, destination, file_size=None, chunk_size=32768):
if file_size is not None:
pbar = tqdm(total=math.ceil(file_size / chunk_size), unit='chunk')
readable_file_size = sizeof_fmt(file_size)
else:
pbar = None
with open(destination, 'wb') as f:
downloaded_size = 0
for chunk in response.iter_content(chunk_size):
downloaded_size += chunk_size
if pbar is not None:
pbar.update(1)
pbar.set_description(f'Download {sizeof_fmt(downloaded_size)} / {readable_file_size}')
if chunk: # filter out keep-alive new chunks
f.write(chunk)
if pbar is not None:
pbar.close()
def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
"""Load file form http url, will download models if necessary.
Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
Args:
url (str): URL to be downloaded.
model_dir (str): The path to save the downloaded model. Should be a full path. If None, use pytorch hub_dir.
Default: None.
progress (bool): Whether to show the download progress. Default: True.
file_name (str): The downloaded file name. If None, use the file name in the url. Default: None.
Returns:
str: The path to the downloaded file.
"""
if model_dir is None: # use the pytorch hub_dir
hub_dir = get_dir()
model_dir = os.path.join(hub_dir, 'checkpoints')
os.makedirs(model_dir, exist_ok=True)
parts = urlparse(url)
filename = os.path.basename(parts.path)
if file_name is not None:
filename = file_name
cached_file = os.path.abspath(os.path.join(model_dir, filename))
if not os.path.exists(cached_file):
print(f'Downloading: "{url}" to {cached_file}\n')
download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
return cached_file

99
preprocessing/matanyone/tools/interact_tools.py Normal file
View File

@@ -0,0 +1,99 @@
import time
import torch
import cv2
from PIL import Image, ImageDraw, ImageOps
import numpy as np
from typing import Union
from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator
import matplotlib.pyplot as plt
import PIL
from .mask_painter import mask_painter as mask_painter2
from .base_segmenter import BaseSegmenter
from .painter import mask_painter, point_painter
import os
import requests
import sys
mask_color = 3
mask_alpha = 0.7
contour_color = 1
contour_width = 5
point_color_ne = 8
point_color_ps = 50
point_alpha = 0.9
point_radius = 15
contour_color = 2
contour_width = 5
class SamControler():
def __init__(self, SAM_checkpoint, model_type, device):
'''
initialize sam controler
'''
self.sam_controler = BaseSegmenter(SAM_checkpoint, model_type, device)
# def seg_again(self, image: np.ndarray):
# '''
# it is used when interact in video
# '''
# self.sam_controler.reset_image()
# self.sam_controler.set_image(image)
# return
def first_frame_click(self, image: np.ndarray, points:np.ndarray, labels: np.ndarray, multimask=True,mask_color=3):
'''
it is used in first frame in video
return: mask, logit, painted image(mask+point)
'''
# self.sam_controler.set_image(image)
origal_image = self.sam_controler.orignal_image
neg_flag = labels[-1]
if neg_flag==1:
#find neg
prompts = {
'point_coords': points,
'point_labels': labels,
}
masks, scores, logits = self.sam_controler.predict(prompts, 'point', multimask)
mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :]
prompts = {
'point_coords': points,
'point_labels': labels,
'mask_input': logit[None, :, :]
}
masks, scores, logits = self.sam_controler.predict(prompts, 'both', multimask)
mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :]
else:
#find positive
prompts = {
'point_coords': points,
'point_labels': labels,
}
masks, scores, logits = self.sam_controler.predict(prompts, 'point', multimask)
mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :]
assert len(points)==len(labels)
painted_image = mask_painter(image, mask.astype('uint8'), mask_color, mask_alpha, contour_color, contour_width)
painted_image = point_painter(painted_image, np.squeeze(points[np.argwhere(labels>0)],axis = 1), point_color_ne, point_alpha, point_radius, contour_color, contour_width)
painted_image = point_painter(painted_image, np.squeeze(points[np.argwhere(labels<1)],axis = 1), point_color_ps, point_alpha, point_radius, contour_color, contour_width)
painted_image = Image.fromarray(painted_image)
return mask, logit, painted_image

288
preprocessing/matanyone/tools/mask_painter.py Normal file
View File

@@ -0,0 +1,288 @@
import cv2
import torch
import numpy as np
from PIL import Image
import copy
import time
def colormap(rgb=True):
color_list = np.array(
[
0.000, 0.000, 0.000,
1.000, 1.000, 1.000,
1.000, 0.498, 0.313,
0.392, 0.581, 0.929,
0.000, 0.447, 0.741,
0.850, 0.325, 0.098,
0.929, 0.694, 0.125,
0.494, 0.184, 0.556,
0.466, 0.674, 0.188,
0.301, 0.745, 0.933,
0.635, 0.078, 0.184,
0.300, 0.300, 0.300,
0.600, 0.600, 0.600,
1.000, 0.000, 0.000,
1.000, 0.500, 0.000,
0.749, 0.749, 0.000,
0.000, 1.000, 0.000,
0.000, 0.000, 1.000,
0.667, 0.000, 1.000,
0.333, 0.333, 0.000,
0.333, 0.667, 0.000,
0.333, 1.000, 0.000,
0.667, 0.333, 0.000,
0.667, 0.667, 0.000,
0.667, 1.000, 0.000,
1.000, 0.333, 0.000,
1.000, 0.667, 0.000,
1.000, 1.000, 0.000,
0.000, 0.333, 0.500,
0.000, 0.667, 0.500,
0.000, 1.000, 0.500,
0.333, 0.000, 0.500,
0.333, 0.333, 0.500,
0.333, 0.667, 0.500,
0.333, 1.000, 0.500,
0.667, 0.000, 0.500,
0.667, 0.333, 0.500,
0.667, 0.667, 0.500,
0.667, 1.000, 0.500,
1.000, 0.000, 0.500,
1.000, 0.333, 0.500,
1.000, 0.667, 0.500,
1.000, 1.000, 0.500,
0.000, 0.333, 1.000,
0.000, 0.667, 1.000,
0.000, 1.000, 1.000,
0.333, 0.000, 1.000,
0.333, 0.333, 1.000,
0.333, 0.667, 1.000,
0.333, 1.000, 1.000,
0.667, 0.000, 1.000,
0.667, 0.333, 1.000,
0.667, 0.667, 1.000,
0.667, 1.000, 1.000,
1.000, 0.000, 1.000,
1.000, 0.333, 1.000,
1.000, 0.667, 1.000,
0.167, 0.000, 0.000,
0.333, 0.000, 0.000,
0.500, 0.000, 0.000,
0.667, 0.000, 0.000,
0.833, 0.000, 0.000,
1.000, 0.000, 0.000,
0.000, 0.167, 0.000,
0.000, 0.333, 0.000,
0.000, 0.500, 0.000,
0.000, 0.667, 0.000,
0.000, 0.833, 0.000,
0.000, 1.000, 0.000,
0.000, 0.000, 0.167,
0.000, 0.000, 0.333,
0.000, 0.000, 0.500,
0.000, 0.000, 0.667,
0.000, 0.000, 0.833,
0.000, 0.000, 1.000,
0.143, 0.143, 0.143,
0.286, 0.286, 0.286,
0.429, 0.429, 0.429,
0.571, 0.571, 0.571,
0.714, 0.714, 0.714,
0.857, 0.857, 0.857
]
).astype(np.float32)
color_list = color_list.reshape((-1, 3)) * 255
if not rgb:
color_list = color_list[:, ::-1]
return color_list
color_list = colormap()
color_list = color_list.astype('uint8').tolist()
def vis_add_mask(image, background_mask, contour_mask, background_color, contour_color, background_alpha, contour_alpha):
background_color = np.array(background_color)
contour_color = np.array(contour_color)
# background_mask = 1 - background_mask
# contour_mask = 1 - contour_mask
for i in range(3):
image[:, :, i] = image[:, :, i] * (1-background_alpha+background_mask*background_alpha) \
+ background_color[i] * (background_alpha-background_mask*background_alpha)
image[:, :, i] = image[:, :, i] * (1-contour_alpha+contour_mask*contour_alpha) \
+ contour_color[i] * (contour_alpha-contour_mask*contour_alpha)
return image.astype('uint8')
def mask_generator_00(mask, background_radius, contour_radius):
# no background width when '00'
# distance map
dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3)
dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3)
dist_map = dist_transform_fore - dist_transform_back
# ...:::!!!:::...
contour_radius += 2
contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
contour_mask = contour_mask / np.max(contour_mask)
contour_mask[contour_mask>0.5] = 1.
return mask, contour_mask
def mask_generator_01(mask, background_radius, contour_radius):
# no background width when '00'
# distance map
dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3)
dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3)
dist_map = dist_transform_fore - dist_transform_back
# ...:::!!!:::...
contour_radius += 2
contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
contour_mask = contour_mask / np.max(contour_mask)
return mask, contour_mask
def mask_generator_10(mask, background_radius, contour_radius):
# distance map
dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3)
dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3)
dist_map = dist_transform_fore - dist_transform_back
# .....:::::!!!!!
background_mask = np.clip(dist_map, -background_radius, background_radius)
background_mask = (background_mask - np.min(background_mask))
background_mask = background_mask / np.max(background_mask)
# ...:::!!!:::...
contour_radius += 2
contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
contour_mask = contour_mask / np.max(contour_mask)
contour_mask[contour_mask>0.5] = 1.
return background_mask, contour_mask
def mask_generator_11(mask, background_radius, contour_radius):
# distance map
dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3)
dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3)
dist_map = dist_transform_fore - dist_transform_back
# .....:::::!!!!!
background_mask = np.clip(dist_map, -background_radius, background_radius)
background_mask = (background_mask - np.min(background_mask))
background_mask = background_mask / np.max(background_mask)
# ...:::!!!:::...
contour_radius += 2
contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
contour_mask = contour_mask / np.max(contour_mask)
return background_mask, contour_mask
def mask_painter(input_image, input_mask, background_alpha=0.5, background_blur_radius=7, contour_width=3, contour_color=3, contour_alpha=1, mode='11'):
"""
Input:
input_image: numpy array
input_mask: numpy array
background_alpha: transparency of background, [0, 1], 1: all black, 0: do nothing
background_blur_radius: radius of background blur, must be odd number
contour_width: width of mask contour, must be odd number
contour_color: color index (in color map) of mask contour, 0: black, 1: white, >1: others
contour_alpha: transparency of mask contour, [0, 1], if 0: no contour highlighted
mode: painting mode, '00', no blur, '01' only blur contour, '10' only blur background, '11' blur both
Output:
painted_image: numpy array
"""
assert input_image.shape[:2] == input_mask.shape, 'different shape'
assert background_blur_radius % 2 * contour_width % 2 > 0, 'background_blur_radius and contour_width must be ODD'
assert mode in ['00', '01', '10', '11'], 'mode should be 00, 01, 10, or 11'
# downsample input image and mask
width, height = input_image.shape[0], input_image.shape[1]
res = 1024
ratio = min(1.0 * res / max(width, height), 1.0)
input_image = cv2.resize(input_image, (int(height*ratio), int(width*ratio)))
input_mask = cv2.resize(input_mask, (int(height*ratio), int(width*ratio)))
# 0: background, 1: foreground
msk = np.clip(input_mask, 0, 1)
# generate masks for background and contour pixels
background_radius = (background_blur_radius - 1) // 2
contour_radius = (contour_width - 1) // 2
generator_dict = {'00':mask_generator_00, '01':mask_generator_01, '10':mask_generator_10, '11':mask_generator_11}
background_mask, contour_mask = generator_dict[mode](msk, background_radius, contour_radius)
# paint
painted_image = vis_add_mask\
(input_image, background_mask, contour_mask, color_list[0], color_list[contour_color], background_alpha, contour_alpha) # black for background
return painted_image
if __name__ == '__main__':
background_alpha = 0.7 # transparency of background 1: all black, 0: do nothing
background_blur_radius = 31 # radius of background blur, must be odd number
contour_width = 11 # contour width, must be odd number
contour_color = 3 # id in color map, 0: black, 1: white, >1: others
contour_alpha = 1 # transparency of background, 0: no contour highlighted
# load input image and mask
input_image = np.array(Image.open('./test_img/painter_input_image.jpg').convert('RGB'))
input_mask = np.array(Image.open('./test_img/painter_input_mask.jpg').convert('P'))
# paint
overall_time_1 = 0
overall_time_2 = 0
overall_time_3 = 0
overall_time_4 = 0
overall_time_5 = 0
for i in range(50):
t2 = time.time()
painted_image_00 = mask_painter(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha, mode='00')
e2 = time.time()
t3 = time.time()
painted_image_10 = mask_painter(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha, mode='10')
e3 = time.time()
t1 = time.time()
painted_image = mask_painter(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha)
e1 = time.time()
t4 = time.time()
painted_image_01 = mask_painter(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha, mode='01')
e4 = time.time()
t5 = time.time()
painted_image_11 = mask_painter(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha, mode='11')
e5 = time.time()
overall_time_1 += (e1 - t1)
overall_time_2 += (e2 - t2)
overall_time_3 += (e3 - t3)
overall_time_4 += (e4 - t4)
overall_time_5 += (e5 - t5)
print(f'average time w gaussian: {overall_time_1/50}')
print(f'average time w/o gaussian00: {overall_time_2/50}')
print(f'average time w/o gaussian10: {overall_time_3/50}')
print(f'average time w/o gaussian01: {overall_time_4/50}')
print(f'average time w/o gaussian11: {overall_time_5/50}')
# save
painted_image_00 = Image.fromarray(painted_image_00)
painted_image_00.save('./test_img/painter_output_image_00.png')
painted_image_10 = Image.fromarray(painted_image_10)
painted_image_10.save('./test_img/painter_output_image_10.png')
painted_image_01 = Image.fromarray(painted_image_01)
painted_image_01.save('./test_img/painter_output_image_01.png')
painted_image_11 = Image.fromarray(painted_image_11)
painted_image_11.save('./test_img/painter_output_image_11.png')

136
preprocessing/matanyone/tools/misc.py Normal file
View File

@@ -0,0 +1,136 @@
import os
import re
import random
import time
import torch
import torch.nn as nn
import logging
import numpy as np
from os import path as osp
def constant_init(module, val, bias=0):
if hasattr(module, 'weight') and module.weight is not None:
nn.init.constant_(module.weight, val)
if hasattr(module, 'bias') and module.bias is not None:
nn.init.constant_(module.bias, bias)
initialized_logger = {}
def get_root_logger(logger_name='basicsr', log_level=logging.INFO, log_file=None):
"""Get the root logger.
The logger will be initialized if it has not been initialized. By default a
StreamHandler will be added. If `log_file` is specified, a FileHandler will
also be added.
Args:
logger_name (str): root logger name. Default: 'basicsr'.
log_file (str | None): The log filename. If specified, a FileHandler
will be added to the root logger.
log_level (int): The root logger level. Note that only the process of
rank 0 is affected, while other processes will set the level to
"Error" and be silent most of the time.
Returns:
logging.Logger: The root logger.
"""
logger = logging.getLogger(logger_name)
# if the logger has been initialized, just return it
if logger_name in initialized_logger:
return logger
format_str = '%(asctime)s %(levelname)s: %(message)s'
stream_handler = logging.StreamHandler()
stream_handler.setFormatter(logging.Formatter(format_str))
logger.addHandler(stream_handler)
logger.propagate = False
if log_file is not None:
logger.setLevel(log_level)
# add file handler
# file_handler = logging.FileHandler(log_file, 'w')
file_handler = logging.FileHandler(log_file, 'a') #Shangchen: keep the previous log
file_handler.setFormatter(logging.Formatter(format_str))
file_handler.setLevel(log_level)
logger.addHandler(file_handler)
initialized_logger[logger_name] = True
return logger
match = re.match(r"^([0-9]+)\.([0-9]+)\.([0-9]+)", torch.__version__)
if match:
version_tuple = match.groups()
IS_HIGH_VERSION = [int(v) for v in version_tuple] >= [1, 12, 0]
else:
logger = get_root_logger()
logger.warning(f"Could not parse torch version '{torch.__version__}'. Assuming it's not a high version >= 1.12.0.")
IS_HIGH_VERSION = False
def gpu_is_available():
if IS_HIGH_VERSION:
if torch.backends.mps.is_available():
return True
return True if torch.cuda.is_available() and torch.backends.cudnn.is_available() else False
def get_device(gpu_id=None):
if gpu_id is None:
gpu_str = ''
elif isinstance(gpu_id, int):
gpu_str = f':{gpu_id}'
else:
raise TypeError('Input should be int value.')
if IS_HIGH_VERSION:
if torch.backends.mps.is_available():
return torch.device('mps'+gpu_str)
return torch.device('cuda'+gpu_str if torch.cuda.is_available() and torch.backends.cudnn.is_available() else 'cpu')
def set_random_seed(seed):
"""Set random seeds."""
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
def get_time_str():
return time.strftime('%Y%m%d_%H%M%S', time.localtime())
def scandir(dir_path, suffix=None, recursive=False, full_path=False):
"""Scan a directory to find the interested files.
Args:
dir_path (str): Path of the directory.
suffix (str | tuple(str), optional): File suffix that we are
interested in. Default: None.
recursive (bool, optional): If set to True, recursively scan the
directory. Default: False.
full_path (bool, optional): If set to True, include the dir_path.
Default: False.
Returns:
A generator for all the interested files with relative pathes.
"""
if (suffix is not None) and not isinstance(suffix, (str, tuple)):
raise TypeError('"suffix" must be a string or tuple of strings')
root = dir_path
def _scandir(dir_path, suffix, recursive):
for entry in os.scandir(dir_path):
if not entry.name.startswith('.') and entry.is_file():
if full_path:
return_path = entry.path
else:
return_path = osp.relpath(entry.path, root)
if suffix is None:
yield return_path
elif return_path.endswith(suffix):
yield return_path
else:
if recursive:
yield from _scandir(entry.path, suffix=suffix, recursive=recursive)
else:
continue
return _scandir(dir_path, suffix=suffix, recursive=recursive)

215
preprocessing/matanyone/tools/painter.py Normal file
View File

@@ -0,0 +1,215 @@
# paint masks, contours, or points on images, with specified colors
import cv2
import torch
import numpy as np
from PIL import Image
import copy
import time
def colormap(rgb=True):
color_list = np.array(
[
0.000, 0.000, 0.000,
1.000, 1.000, 1.000,
1.000, 0.498, 0.313,
0.392, 0.581, 0.929,
0.000, 0.447, 0.741,
0.850, 0.325, 0.098,
0.929, 0.694, 0.125,
0.494, 0.184, 0.556,
0.466, 0.674, 0.188,
0.301, 0.745, 0.933,
0.635, 0.078, 0.184,
0.300, 0.300, 0.300,
0.600, 0.600, 0.600,
1.000, 0.000, 0.000,
1.000, 0.500, 0.000,
0.749, 0.749, 0.000,
0.000, 1.000, 0.000,
0.000, 0.000, 1.000,
0.667, 0.000, 1.000,
0.333, 0.333, 0.000,
0.333, 0.667, 0.000,
0.333, 1.000, 0.000,
0.667, 0.333, 0.000,
0.667, 0.667, 0.000,
0.667, 1.000, 0.000,
1.000, 0.333, 0.000,
1.000, 0.667, 0.000,
1.000, 1.000, 0.000,
0.000, 0.333, 0.500,
0.000, 0.667, 0.500,
0.000, 1.000, 0.500,
0.333, 0.000, 0.500,
0.333, 0.333, 0.500,
0.333, 0.667, 0.500,
0.333, 1.000, 0.500,
0.667, 0.000, 0.500,
0.667, 0.333, 0.500,
0.667, 0.667, 0.500,
0.667, 1.000, 0.500,
1.000, 0.000, 0.500,
1.000, 0.333, 0.500,
1.000, 0.667, 0.500,
1.000, 1.000, 0.500,
0.000, 0.333, 1.000,
0.000, 0.667, 1.000,
0.000, 1.000, 1.000,
0.333, 0.000, 1.000,
0.333, 0.333, 1.000,
0.333, 0.667, 1.000,
0.333, 1.000, 1.000,
0.667, 0.000, 1.000,
0.667, 0.333, 1.000,
0.667, 0.667, 1.000,
0.667, 1.000, 1.000,
1.000, 0.000, 1.000,
1.000, 0.333, 1.000,
1.000, 0.667, 1.000,
0.167, 0.000, 0.000,
0.333, 0.000, 0.000,
0.500, 0.000, 0.000,
0.667, 0.000, 0.000,
0.833, 0.000, 0.000,
1.000, 0.000, 0.000,
0.000, 0.167, 0.000,
0.000, 0.333, 0.000,
0.000, 0.500, 0.000,
0.000, 0.667, 0.000,
0.000, 0.833, 0.000,
0.000, 1.000, 0.000,
0.000, 0.000, 0.167,
0.000, 0.000, 0.333,
0.000, 0.000, 0.500,
0.000, 0.000, 0.667,
0.000, 0.000, 0.833,
0.000, 0.000, 1.000,
0.143, 0.143, 0.143,
0.286, 0.286, 0.286,
0.429, 0.429, 0.429,
0.571, 0.571, 0.571,
0.714, 0.714, 0.714,
0.857, 0.857, 0.857
]
).astype(np.float32)
color_list = color_list.reshape((-1, 3)) * 255
if not rgb:
color_list = color_list[:, ::-1]
return color_list
color_list = colormap()
color_list = color_list.astype('uint8').tolist()
def vis_add_mask(image, mask, color, alpha):
color = np.array(color_list[color])
mask = mask > 0.5
image[mask] = image[mask] * (1-alpha) + color * alpha
return image.astype('uint8')
def point_painter(input_image, input_points, point_color=5, point_alpha=0.9, point_radius=15, contour_color=2, contour_width=5):
h, w = input_image.shape[:2]
point_mask = np.zeros((h, w)).astype('uint8')
for point in input_points:
point_mask[point[1], point[0]] = 1
kernel = cv2.getStructuringElement(2, (point_radius, point_radius))
point_mask = cv2.dilate(point_mask, kernel)
contour_radius = (contour_width - 1) // 2
dist_transform_fore = cv2.distanceTransform(point_mask, cv2.DIST_L2, 3)
dist_transform_back = cv2.distanceTransform(1-point_mask, cv2.DIST_L2, 3)
dist_map = dist_transform_fore - dist_transform_back
# ...:::!!!:::...
contour_radius += 2
contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
contour_mask = contour_mask / np.max(contour_mask)
contour_mask[contour_mask>0.5] = 1.
# paint mask
painted_image = vis_add_mask(input_image.copy(), point_mask, point_color, point_alpha)
# paint contour
painted_image = vis_add_mask(painted_image.copy(), 1-contour_mask, contour_color, 1)
return painted_image
def mask_painter(input_image, input_mask, mask_color=5, mask_alpha=0.7, contour_color=1, contour_width=3):
assert input_image.shape[:2] == input_mask.shape, 'different shape between image and mask'
# 0: background, 1: foreground
mask = np.clip(input_mask, 0, 1)
contour_radius = (contour_width - 1) // 2
dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3)
dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3)
dist_map = dist_transform_fore - dist_transform_back
# ...:::!!!:::...
contour_radius += 2
contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
contour_mask = contour_mask / np.max(contour_mask)
contour_mask[contour_mask>0.5] = 1.
# paint mask
painted_image = vis_add_mask(input_image.copy(), mask.copy(), mask_color, mask_alpha)
# paint contour
painted_image = vis_add_mask(painted_image.copy(), 1-contour_mask, contour_color, 1)
return painted_image
def background_remover(input_image, input_mask):
"""
input_image: H, W, 3, np.array
input_mask: H, W, np.array
image_wo_background: PIL.Image
"""
assert input_image.shape[:2] == input_mask.shape, 'different shape between image and mask'
# 0: background, 1: foreground
mask = np.expand_dims(np.clip(input_mask, 0, 1), axis=2)*255
image_wo_background = np.concatenate([input_image, mask], axis=2) # H, W, 4
image_wo_background = Image.fromarray(image_wo_background).convert('RGBA')
return image_wo_background
if __name__ == '__main__':
input_image = np.array(Image.open('images/painter_input_image.jpg').convert('RGB'))
input_mask = np.array(Image.open('images/painter_input_mask.jpg').convert('P'))
# example of mask painter
mask_color = 3
mask_alpha = 0.7
contour_color = 1
contour_width = 5
# save
painted_image = Image.fromarray(input_image)
painted_image.save('images/original.png')
painted_image = mask_painter(input_image, input_mask, mask_color, mask_alpha, contour_color, contour_width)
# save
painted_image = Image.fromarray(input_image)
painted_image.save('images/original1.png')
# example of point painter
input_image = np.array(Image.open('images/painter_input_image.jpg').convert('RGB'))
input_points = np.array([[500, 375], [70, 600]]) # x, y
point_color = 5
point_alpha = 0.9
point_radius = 15
contour_color = 2
contour_width = 5
painted_image_1 = point_painter(input_image, input_points, point_color, point_alpha, point_radius, contour_color, contour_width)
# save
painted_image = Image.fromarray(painted_image_1)
painted_image.save('images/point_painter_1.png')
input_image = np.array(Image.open('images/painter_input_image.jpg').convert('RGB'))
painted_image_2 = point_painter(input_image, input_points, point_color=9, point_radius=20, contour_color=29)
# save
painted_image = Image.fromarray(painted_image_2)
painted_image.save('images/point_painter_2.png')
# example of background remover
input_image = np.array(Image.open('images/original.png').convert('RGB'))
image_wo_background = background_remover(input_image, input_mask) # return PIL.Image
image_wo_background.save('images/image_wo_background.png')

0
preprocessing/matanyone/utils/__init__.py Normal file
View File

27
preprocessing/matanyone/utils/get_default_model.py Normal file
View File

@@ -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 ..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

62
preprocessing/matanyone/utils/tensor_utils.py Normal file
View File

@@ -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"):
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

2
preprocessing/midas/__init__.py Normal file
View File

@@ -0,0 +1,2 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.

166
preprocessing/midas/api.py Normal file
View File

@@ -0,0 +1,166 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
# based on https://github.com/isl-org/MiDaS
import cv2
import torch
import torch.nn as nn
from torchvision.transforms import Compose
from .dpt_depth import DPTDepthModel
from .midas_net import MidasNet
from .midas_net_custom import MidasNet_small
from .transforms import NormalizeImage, PrepareForNet, Resize
# ISL_PATHS = {
# "dpt_large": "dpt_large-midas-2f21e586.pt",
# "dpt_hybrid": "dpt_hybrid-midas-501f0c75.pt",
# "midas_v21": "",
# "midas_v21_small": "",
# }
# remote_model_path =
# "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt"
def disabled_train(self, mode=True):
"""Overwrite model.train with this function to make sure train/eval mode
does not change anymore."""
return self
def load_midas_transform(model_type):
# https://github.com/isl-org/MiDaS/blob/master/run.py
# load transform only
if model_type == 'dpt_large': # DPT-Large
net_w, net_h = 384, 384
resize_mode = 'minimal'
normalization = NormalizeImage(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
elif model_type == 'dpt_hybrid': # DPT-Hybrid
net_w, net_h = 384, 384
resize_mode = 'minimal'
normalization = NormalizeImage(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
elif model_type == 'midas_v21':
net_w, net_h = 384, 384
resize_mode = 'upper_bound'
normalization = NormalizeImage(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
elif model_type == 'midas_v21_small':
net_w, net_h = 256, 256
resize_mode = 'upper_bound'
normalization = NormalizeImage(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
else:
assert False, f"model_type '{model_type}' not implemented, use: --model_type large"
transform = Compose([
Resize(
net_w,
net_h,
resize_target=None,
keep_aspect_ratio=True,
ensure_multiple_of=32,
resize_method=resize_mode,
image_interpolation_method=cv2.INTER_CUBIC,
),
normalization,
PrepareForNet(),
])
return transform
def load_model(model_type, model_path):
# https://github.com/isl-org/MiDaS/blob/master/run.py
# load network
# model_path = ISL_PATHS[model_type]
if model_type == 'dpt_large': # DPT-Large
model = DPTDepthModel(
path=model_path,
backbone='vitl16_384',
non_negative=True,
)
net_w, net_h = 384, 384
resize_mode = 'minimal'
normalization = NormalizeImage(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
elif model_type == 'dpt_hybrid': # DPT-Hybrid
model = DPTDepthModel(
path=model_path,
backbone='vitb_rn50_384',
non_negative=True,
)
net_w, net_h = 384, 384
resize_mode = 'minimal'
normalization = NormalizeImage(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
elif model_type == 'midas_v21':
model = MidasNet(model_path, non_negative=True)
net_w, net_h = 384, 384
resize_mode = 'upper_bound'
normalization = NormalizeImage(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
elif model_type == 'midas_v21_small':
model = MidasNet_small(model_path,
features=64,
backbone='efficientnet_lite3',
exportable=True,
non_negative=True,
blocks={'expand': True})
net_w, net_h = 256, 256
resize_mode = 'upper_bound'
normalization = NormalizeImage(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
else:
print(
f"model_type '{model_type}' not implemented, use: --model_type large"
)
assert False
transform = Compose([
Resize(
net_w,
net_h,
resize_target=None,
keep_aspect_ratio=True,
ensure_multiple_of=32,
resize_method=resize_mode,
image_interpolation_method=cv2.INTER_CUBIC,
),
normalization,
PrepareForNet(),
])
return model.eval(), transform
class MiDaSInference(nn.Module):
MODEL_TYPES_TORCH_HUB = ['DPT_Large', 'DPT_Hybrid', 'MiDaS_small']
MODEL_TYPES_ISL = [
'dpt_large',
'dpt_hybrid',
'midas_v21',
'midas_v21_small',
]
def __init__(self, model_type, model_path):
super().__init__()
assert (model_type in self.MODEL_TYPES_ISL)
model, _ = load_model(model_type, model_path)
self.model = model
self.model.train = disabled_train
def forward(self, x):
with torch.no_grad():
prediction = self.model(x)
return prediction

18
preprocessing/midas/base_model.py Normal file
View File

@@ -0,0 +1,18 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import torch
class BaseModel(torch.nn.Module):
def load(self, path):
"""Load model from file.
Args:
path (str): file path
"""
parameters = torch.load(path, map_location=torch.device('cpu'), weights_only=True)
if 'optimizer' in parameters:
parameters = parameters['model']
self.load_state_dict(parameters)

391
preprocessing/midas/blocks.py Normal file
View File

@@ -0,0 +1,391 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import torch
import torch.nn as nn
from .vit import (_make_pretrained_vitb16_384, _make_pretrained_vitb_rn50_384,
_make_pretrained_vitl16_384)
def _make_encoder(
backbone,
features,
use_pretrained,
groups=1,
expand=False,
exportable=True,
hooks=None,
use_vit_only=False,
use_readout='ignore',
):
if backbone == 'vitl16_384':
pretrained = _make_pretrained_vitl16_384(use_pretrained,
hooks=hooks,
use_readout=use_readout)
scratch = _make_scratch(
[256, 512, 1024, 1024], features, groups=groups,
expand=expand) # ViT-L/16 - 85.0% Top1 (backbone)
elif backbone == 'vitb_rn50_384':
pretrained = _make_pretrained_vitb_rn50_384(
use_pretrained,
hooks=hooks,
use_vit_only=use_vit_only,
use_readout=use_readout,
)
scratch = _make_scratch(
[256, 512, 768, 768], features, groups=groups,
expand=expand) # ViT-H/16 - 85.0% Top1 (backbone)
elif backbone == 'vitb16_384':
pretrained = _make_pretrained_vitb16_384(use_pretrained,
hooks=hooks,
use_readout=use_readout)
scratch = _make_scratch(
[96, 192, 384, 768], features, groups=groups,
expand=expand) # ViT-B/16 - 84.6% Top1 (backbone)
elif backbone == 'resnext101_wsl':
pretrained = _make_pretrained_resnext101_wsl(use_pretrained)
scratch = _make_scratch([256, 512, 1024, 2048],
features,
groups=groups,
expand=expand) # efficientnet_lite3
elif backbone == 'efficientnet_lite3':
pretrained = _make_pretrained_efficientnet_lite3(use_pretrained,
exportable=exportable)
scratch = _make_scratch([32, 48, 136, 384],
features,
groups=groups,
expand=expand) # efficientnet_lite3
else:
print(f"Backbone '{backbone}' not implemented")
assert False
return pretrained, scratch
def _make_scratch(in_shape, out_shape, groups=1, expand=False):
scratch = nn.Module()
out_shape1 = out_shape
out_shape2 = out_shape
out_shape3 = out_shape
out_shape4 = out_shape
if expand is True:
out_shape1 = out_shape
out_shape2 = out_shape * 2
out_shape3 = out_shape * 4
out_shape4 = out_shape * 8
scratch.layer1_rn = nn.Conv2d(in_shape[0],
out_shape1,
kernel_size=3,
stride=1,
padding=1,
bias=False,
groups=groups)
scratch.layer2_rn = nn.Conv2d(in_shape[1],
out_shape2,
kernel_size=3,
stride=1,
padding=1,
bias=False,
groups=groups)
scratch.layer3_rn = nn.Conv2d(in_shape[2],
out_shape3,
kernel_size=3,
stride=1,
padding=1,
bias=False,
groups=groups)
scratch.layer4_rn = nn.Conv2d(in_shape[3],
out_shape4,
kernel_size=3,
stride=1,
padding=1,
bias=False,
groups=groups)
return scratch
def _make_pretrained_efficientnet_lite3(use_pretrained, exportable=False):
efficientnet = torch.hub.load('rwightman/gen-efficientnet-pytorch',
'tf_efficientnet_lite3',
pretrained=use_pretrained,
exportable=exportable)
return _make_efficientnet_backbone(efficientnet)
def _make_efficientnet_backbone(effnet):
pretrained = nn.Module()
pretrained.layer1 = nn.Sequential(effnet.conv_stem, effnet.bn1,
effnet.act1, *effnet.blocks[0:2])
pretrained.layer2 = nn.Sequential(*effnet.blocks[2:3])
pretrained.layer3 = nn.Sequential(*effnet.blocks[3:5])
pretrained.layer4 = nn.Sequential(*effnet.blocks[5:9])
return pretrained
def _make_resnet_backbone(resnet):
pretrained = nn.Module()
pretrained.layer1 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu,
resnet.maxpool, resnet.layer1)
pretrained.layer2 = resnet.layer2
pretrained.layer3 = resnet.layer3
pretrained.layer4 = resnet.layer4
return pretrained
def _make_pretrained_resnext101_wsl(use_pretrained):
resnet = torch.hub.load('facebookresearch/WSL-Images',
'resnext101_32x8d_wsl')
return _make_resnet_backbone(resnet)
class Interpolate(nn.Module):
"""Interpolation module.
"""
def __init__(self, scale_factor, mode, align_corners=False):
"""Init.
Args:
scale_factor (float): scaling
mode (str): interpolation mode
"""
super(Interpolate, self).__init__()
self.interp = nn.functional.interpolate
self.scale_factor = scale_factor
self.mode = mode
self.align_corners = align_corners
def forward(self, x):
"""Forward pass.
Args:
x (tensor): input
Returns:
tensor: interpolated data
"""
x = self.interp(x,
scale_factor=self.scale_factor,
mode=self.mode,
align_corners=self.align_corners)
return x
class ResidualConvUnit(nn.Module):
"""Residual convolution module.
"""
def __init__(self, features):
"""Init.
Args:
features (int): number of features
"""
super().__init__()
self.conv1 = nn.Conv2d(features,
features,
kernel_size=3,
stride=1,
padding=1,
bias=True)
self.conv2 = nn.Conv2d(features,
features,
kernel_size=3,
stride=1,
padding=1,
bias=True)
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
"""Forward pass.
Args:
x (tensor): input
Returns:
tensor: output
"""
out = self.relu(x)
out = self.conv1(out)
out = self.relu(out)
out = self.conv2(out)
return out + x
class FeatureFusionBlock(nn.Module):
"""Feature fusion block.
"""
def __init__(self, features):
"""Init.
Args:
features (int): number of features
"""
super(FeatureFusionBlock, self).__init__()
self.resConfUnit1 = ResidualConvUnit(features)
self.resConfUnit2 = ResidualConvUnit(features)
def forward(self, *xs):
"""Forward pass.
Returns:
tensor: output
"""
output = xs[0]
if len(xs) == 2:
output += self.resConfUnit1(xs[1])
output = self.resConfUnit2(output)
output = nn.functional.interpolate(output,
scale_factor=2,
mode='bilinear',
align_corners=True)
return output
class ResidualConvUnit_custom(nn.Module):
"""Residual convolution module.
"""
def __init__(self, features, activation, bn):
"""Init.
Args:
features (int): number of features
"""
super().__init__()
self.bn = bn
self.groups = 1
self.conv1 = nn.Conv2d(features,
features,
kernel_size=3,
stride=1,
padding=1,
bias=True,
groups=self.groups)
self.conv2 = nn.Conv2d(features,
features,
kernel_size=3,
stride=1,
padding=1,
bias=True,
groups=self.groups)
if self.bn is True:
self.bn1 = nn.BatchNorm2d(features)
self.bn2 = nn.BatchNorm2d(features)
self.activation = activation
self.skip_add = nn.quantized.FloatFunctional()
def forward(self, x):
"""Forward pass.
Args:
x (tensor): input
Returns:
tensor: output
"""
out = self.activation(x)
out = self.conv1(out)
if self.bn is True:
out = self.bn1(out)
out = self.activation(out)
out = self.conv2(out)
if self.bn is True:
out = self.bn2(out)
if self.groups > 1:
out = self.conv_merge(out)
return self.skip_add.add(out, x)
# return out + x
class FeatureFusionBlock_custom(nn.Module):
"""Feature fusion block.
"""
def __init__(self,
features,
activation,
deconv=False,
bn=False,
expand=False,
align_corners=True):
"""Init.
Args:
features (int): number of features
"""
super(FeatureFusionBlock_custom, self).__init__()
self.deconv = deconv
self.align_corners = align_corners
self.groups = 1
self.expand = expand
out_features = features
if self.expand is True:
out_features = features // 2
self.out_conv = nn.Conv2d(features,
out_features,
kernel_size=1,
stride=1,
padding=0,
bias=True,
groups=1)
self.resConfUnit1 = ResidualConvUnit_custom(features, activation, bn)
self.resConfUnit2 = ResidualConvUnit_custom(features, activation, bn)
self.skip_add = nn.quantized.FloatFunctional()
def forward(self, *xs):
"""Forward pass.
Returns:
tensor: output
"""
output = xs[0]
if len(xs) == 2:
res = self.resConfUnit1(xs[1])
output = self.skip_add.add(output, res)
# output += res
output = self.resConfUnit2(output)
output = nn.functional.interpolate(output,
scale_factor=2,
mode='bilinear',
align_corners=self.align_corners)
output = self.out_conv(output)
return output

84
preprocessing/midas/depth.py Normal file
View File

@@ -0,0 +1,84 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import numpy as np
import torch
from einops import rearrange
from PIL import Image
import cv2
def convert_to_numpy(image):
if isinstance(image, Image.Image):
image = np.array(image)
elif isinstance(image, torch.Tensor):
image = image.detach().cpu().numpy()
elif isinstance(image, np.ndarray):
image = image.copy()
else:
raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
return image
def resize_image(input_image, resolution):
H, W, C = input_image.shape
H = float(H)
W = float(W)
k = float(resolution) / min(H, W)
H *= k
W *= k
H = int(np.round(H / 64.0)) * 64
W = int(np.round(W / 64.0)) * 64
img = cv2.resize(
input_image, (W, H),
interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
return img, k
def resize_image_ori(h, w, image, k):
img = cv2.resize(
image, (w, h),
interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
return img
class DepthAnnotator:
def __init__(self, cfg, device=None):
from .api import MiDaSInference
pretrained_model = cfg['PRETRAINED_MODEL']
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if device is None else device
self.model = MiDaSInference(model_type='dpt_hybrid', model_path=pretrained_model).to(self.device)
self.a = cfg.get('A', np.pi * 2.0)
self.bg_th = cfg.get('BG_TH', 0.1)
@torch.no_grad()
@torch.inference_mode()
@torch.autocast('cuda', enabled=False)
def forward(self, image):
image = convert_to_numpy(image)
image_depth = image
h, w, c = image.shape
image_depth, k = resize_image(image_depth,
1024 if min(h, w) > 1024 else min(h, w))
image_depth = torch.from_numpy(image_depth).float().to(self.device)
image_depth = image_depth / 127.5 - 1.0
image_depth = rearrange(image_depth, 'h w c -> 1 c h w')
depth = self.model(image_depth)[0]
depth_pt = depth.clone()
depth_pt -= torch.min(depth_pt)
depth_pt /= torch.max(depth_pt)
depth_pt = depth_pt.cpu().numpy()
depth_image = (depth_pt * 255.0).clip(0, 255).astype(np.uint8)
depth_image = depth_image[..., None].repeat(3, 2)
depth_image = resize_image_ori(h, w, depth_image, k)
return depth_image
class DepthVideoAnnotator(DepthAnnotator):
def forward(self, frames):
ret_frames = []
for frame in frames:
anno_frame = super().forward(np.array(frame))
ret_frames.append(anno_frame)
return ret_frames

107
preprocessing/midas/dpt_depth.py Normal file
View File

@@ -0,0 +1,107 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import torch
import torch.nn as nn
from .base_model import BaseModel
from .blocks import FeatureFusionBlock_custom, Interpolate, _make_encoder
from .vit import forward_vit
def _make_fusion_block(features, use_bn):
return FeatureFusionBlock_custom(
features,
nn.ReLU(False),
deconv=False,
bn=use_bn,
expand=False,
align_corners=True,
)
class DPT(BaseModel):
def __init__(
self,
head,
features=256,
backbone='vitb_rn50_384',
readout='project',
channels_last=False,
use_bn=False,
):
super(DPT, self).__init__()
self.channels_last = channels_last
hooks = {
'vitb_rn50_384': [0, 1, 8, 11],
'vitb16_384': [2, 5, 8, 11],
'vitl16_384': [5, 11, 17, 23],
}
# Instantiate backbone and reassemble blocks
self.pretrained, self.scratch = _make_encoder(
backbone,
features,
False, # Set to true of you want to train from scratch, uses ImageNet weights
groups=1,
expand=False,
exportable=False,
hooks=hooks[backbone],
use_readout=readout,
)
self.scratch.refinenet1 = _make_fusion_block(features, use_bn)
self.scratch.refinenet2 = _make_fusion_block(features, use_bn)
self.scratch.refinenet3 = _make_fusion_block(features, use_bn)
self.scratch.refinenet4 = _make_fusion_block(features, use_bn)
self.scratch.output_conv = head
def forward(self, x):
if self.channels_last is True:
x.contiguous(memory_format=torch.channels_last)
layer_1, layer_2, layer_3, layer_4 = forward_vit(self.pretrained, x)
layer_1_rn = self.scratch.layer1_rn(layer_1)
layer_2_rn = self.scratch.layer2_rn(layer_2)
layer_3_rn = self.scratch.layer3_rn(layer_3)
layer_4_rn = self.scratch.layer4_rn(layer_4)
path_4 = self.scratch.refinenet4(layer_4_rn)
path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
out = self.scratch.output_conv(path_1)
return out
class DPTDepthModel(DPT):
def __init__(self, path=None, non_negative=True, **kwargs):
features = kwargs['features'] if 'features' in kwargs else 256
head = nn.Sequential(
nn.Conv2d(features,
features // 2,
kernel_size=3,
stride=1,
padding=1),
Interpolate(scale_factor=2, mode='bilinear', align_corners=True),
nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1),
nn.ReLU(True),
nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
nn.ReLU(True) if non_negative else nn.Identity(),
nn.Identity(),
)
super().__init__(head, **kwargs)
if path is not None:
self.load(path)
def forward(self, x):
return super().forward(x).squeeze(dim=1)

80
preprocessing/midas/midas_net.py Normal file
View File

@@ -0,0 +1,80 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
"""MidashNet: Network for monocular depth estimation trained by mixing several datasets.
This file contains code that is adapted from
https://github.com/thomasjpfan/pytorch_refinenet/blob/master/pytorch_refinenet/refinenet/refinenet_4cascade.py
"""
import torch
import torch.nn as nn
from .base_model import BaseModel
from .blocks import FeatureFusionBlock, Interpolate, _make_encoder
class MidasNet(BaseModel):
"""Network for monocular depth estimation.
"""
def __init__(self, path=None, features=256, non_negative=True):
"""Init.
Args:
path (str, optional): Path to saved model. Defaults to None.
features (int, optional): Number of features. Defaults to 256.
backbone (str, optional): Backbone network for encoder. Defaults to resnet50
"""
print('Loading weights: ', path)
super(MidasNet, self).__init__()
use_pretrained = False if path is None else True
self.pretrained, self.scratch = _make_encoder(
backbone='resnext101_wsl',
features=features,
use_pretrained=use_pretrained)
self.scratch.refinenet4 = FeatureFusionBlock(features)
self.scratch.refinenet3 = FeatureFusionBlock(features)
self.scratch.refinenet2 = FeatureFusionBlock(features)
self.scratch.refinenet1 = FeatureFusionBlock(features)
self.scratch.output_conv = nn.Sequential(
nn.Conv2d(features, 128, kernel_size=3, stride=1, padding=1),
Interpolate(scale_factor=2, mode='bilinear'),
nn.Conv2d(128, 32, kernel_size=3, stride=1, padding=1),
nn.ReLU(True),
nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
nn.ReLU(True) if non_negative else nn.Identity(),
)
if path:
self.load(path)
def forward(self, x):
"""Forward pass.
Args:
x (tensor): input data (image)
Returns:
tensor: depth
"""
layer_1 = self.pretrained.layer1(x)
layer_2 = self.pretrained.layer2(layer_1)
layer_3 = self.pretrained.layer3(layer_2)
layer_4 = self.pretrained.layer4(layer_3)
layer_1_rn = self.scratch.layer1_rn(layer_1)
layer_2_rn = self.scratch.layer2_rn(layer_2)
layer_3_rn = self.scratch.layer3_rn(layer_3)
layer_4_rn = self.scratch.layer4_rn(layer_4)
path_4 = self.scratch.refinenet4(layer_4_rn)
path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
out = self.scratch.output_conv(path_1)
return torch.squeeze(out, dim=1)

167
preprocessing/midas/midas_net_custom.py Normal file
View File

@@ -0,0 +1,167 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
"""MidashNet: Network for monocular depth estimation trained by mixing several datasets.
This file contains code that is adapted from
https://github.com/thomasjpfan/pytorch_refinenet/blob/master/pytorch_refinenet/refinenet/refinenet_4cascade.py
"""
import torch
import torch.nn as nn
from .base_model import BaseModel
from .blocks import FeatureFusionBlock_custom, Interpolate, _make_encoder
class MidasNet_small(BaseModel):
"""Network for monocular depth estimation.
"""
def __init__(self,
path=None,
features=64,
backbone='efficientnet_lite3',
non_negative=True,
exportable=True,
channels_last=False,
align_corners=True,
blocks={'expand': True}):
"""Init.
Args:
path (str, optional): Path to saved model. Defaults to None.
features (int, optional): Number of features. Defaults to 256.
backbone (str, optional): Backbone network for encoder. Defaults to resnet50
"""
print('Loading weights: ', path)
super(MidasNet_small, self).__init__()
use_pretrained = False if path else True
self.channels_last = channels_last
self.blocks = blocks
self.backbone = backbone
self.groups = 1
features1 = features
features2 = features
features3 = features
features4 = features
self.expand = False
if 'expand' in self.blocks and self.blocks['expand'] is True:
self.expand = True
features1 = features
features2 = features * 2
features3 = features * 4
features4 = features * 8
self.pretrained, self.scratch = _make_encoder(self.backbone,
features,
use_pretrained,
groups=self.groups,
expand=self.expand,
exportable=exportable)
self.scratch.activation = nn.ReLU(False)
self.scratch.refinenet4 = FeatureFusionBlock_custom(
features4,
self.scratch.activation,
deconv=False,
bn=False,
expand=self.expand,
align_corners=align_corners)
self.scratch.refinenet3 = FeatureFusionBlock_custom(
features3,
self.scratch.activation,
deconv=False,
bn=False,
expand=self.expand,
align_corners=align_corners)
self.scratch.refinenet2 = FeatureFusionBlock_custom(
features2,
self.scratch.activation,
deconv=False,
bn=False,
expand=self.expand,
align_corners=align_corners)
self.scratch.refinenet1 = FeatureFusionBlock_custom(
features1,
self.scratch.activation,
deconv=False,
bn=False,
align_corners=align_corners)
self.scratch.output_conv = nn.Sequential(
nn.Conv2d(features,
features // 2,
kernel_size=3,
stride=1,
padding=1,
groups=self.groups),
Interpolate(scale_factor=2, mode='bilinear'),
nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1),
self.scratch.activation,
nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
nn.ReLU(True) if non_negative else nn.Identity(),
nn.Identity(),
)
if path:
self.load(path)
def forward(self, x):
"""Forward pass.
Args:
x (tensor): input data (image)
Returns:
tensor: depth
"""
if self.channels_last is True:
print('self.channels_last = ', self.channels_last)
x.contiguous(memory_format=torch.channels_last)
layer_1 = self.pretrained.layer1(x)
layer_2 = self.pretrained.layer2(layer_1)
layer_3 = self.pretrained.layer3(layer_2)
layer_4 = self.pretrained.layer4(layer_3)
layer_1_rn = self.scratch.layer1_rn(layer_1)
layer_2_rn = self.scratch.layer2_rn(layer_2)
layer_3_rn = self.scratch.layer3_rn(layer_3)
layer_4_rn = self.scratch.layer4_rn(layer_4)
path_4 = self.scratch.refinenet4(layer_4_rn)
path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
out = self.scratch.output_conv(path_1)
return torch.squeeze(out, dim=1)
def fuse_model(m):
prev_previous_type = nn.Identity()
prev_previous_name = ''
previous_type = nn.Identity()
previous_name = ''
for name, module in m.named_modules():
if prev_previous_type == nn.Conv2d and previous_type == nn.BatchNorm2d and type(
module) == nn.ReLU:
# print("FUSED ", prev_previous_name, previous_name, name)
torch.quantization.fuse_modules(
m, [prev_previous_name, previous_name, name], inplace=True)
elif prev_previous_type == nn.Conv2d and previous_type == nn.BatchNorm2d:
# print("FUSED ", prev_previous_name, previous_name)
torch.quantization.fuse_modules(
m, [prev_previous_name, previous_name], inplace=True)
# elif previous_type == nn.Conv2d and type(module) == nn.ReLU:
# print("FUSED ", previous_name, name)
# torch.quantization.fuse_modules(m, [previous_name, name], inplace=True)
prev_previous_type = previous_type
prev_previous_name = previous_name
previous_type = type(module)
previous_name = name

231
preprocessing/midas/transforms.py Normal file
View File

@@ -0,0 +1,231 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import math
import cv2
import numpy as np
def apply_min_size(sample, size, image_interpolation_method=cv2.INTER_AREA):
"""Rezise the sample to ensure the given size. Keeps aspect ratio.
Args:
sample (dict): sample
size (tuple): image size
Returns:
tuple: new size
"""
shape = list(sample['disparity'].shape)
if shape[0] >= size[0] and shape[1] >= size[1]:
return sample
scale = [0, 0]
scale[0] = size[0] / shape[0]
scale[1] = size[1] / shape[1]
scale = max(scale)
shape[0] = math.ceil(scale * shape[0])
shape[1] = math.ceil(scale * shape[1])
# resize
sample['image'] = cv2.resize(sample['image'],
tuple(shape[::-1]),
interpolation=image_interpolation_method)
sample['disparity'] = cv2.resize(sample['disparity'],
tuple(shape[::-1]),
interpolation=cv2.INTER_NEAREST)
sample['mask'] = cv2.resize(
sample['mask'].astype(np.float32),
tuple(shape[::-1]),
interpolation=cv2.INTER_NEAREST,
)
sample['mask'] = sample['mask'].astype(bool)
return tuple(shape)
class Resize(object):
"""Resize sample to given size (width, height).
"""
def __init__(
self,
width,
height,
resize_target=True,
keep_aspect_ratio=False,
ensure_multiple_of=1,
resize_method='lower_bound',
image_interpolation_method=cv2.INTER_AREA,
):
"""Init.
Args:
width (int): desired output width
height (int): desired output height
resize_target (bool, optional):
True: Resize the full sample (image, mask, target).
False: Resize image only.
Defaults to True.
keep_aspect_ratio (bool, optional):
True: Keep the aspect ratio of the input sample.
Output sample might not have the given width and height, and
resize behaviour depends on the parameter 'resize_method'.
Defaults to False.
ensure_multiple_of (int, optional):
Output width and height is constrained to be multiple of this parameter.
Defaults to 1.
resize_method (str, optional):
"lower_bound": Output will be at least as large as the given size.
"upper_bound": Output will be at max as large as the given size. "
"(Output size might be smaller than given size.)"
"minimal": Scale as least as possible. (Output size might be smaller than given size.)
Defaults to "lower_bound".
"""
self.__width = width
self.__height = height
self.__resize_target = resize_target
self.__keep_aspect_ratio = keep_aspect_ratio
self.__multiple_of = ensure_multiple_of
self.__resize_method = resize_method
self.__image_interpolation_method = image_interpolation_method
def constrain_to_multiple_of(self, x, min_val=0, max_val=None):
y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
if max_val is not None and y > max_val:
y = (np.floor(x / self.__multiple_of) *
self.__multiple_of).astype(int)
if y < min_val:
y = (np.ceil(x / self.__multiple_of) *
self.__multiple_of).astype(int)
return y
def get_size(self, width, height):
# determine new height and width
scale_height = self.__height / height
scale_width = self.__width / width
if self.__keep_aspect_ratio:
if self.__resize_method == 'lower_bound':
# scale such that output size is lower bound
if scale_width > scale_height:
# fit width
scale_height = scale_width
else:
# fit height
scale_width = scale_height
elif self.__resize_method == 'upper_bound':
# scale such that output size is upper bound
if scale_width < scale_height:
# fit width
scale_height = scale_width
else:
# fit height
scale_width = scale_height
elif self.__resize_method == 'minimal':
# scale as least as possbile
if abs(1 - scale_width) < abs(1 - scale_height):
# fit width
scale_height = scale_width
else:
# fit height
scale_width = scale_height
else:
raise ValueError(
f'resize_method {self.__resize_method} not implemented')
if self.__resize_method == 'lower_bound':
new_height = self.constrain_to_multiple_of(scale_height * height,
min_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width,
min_val=self.__width)
elif self.__resize_method == 'upper_bound':
new_height = self.constrain_to_multiple_of(scale_height * height,
max_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width,
max_val=self.__width)
elif self.__resize_method == 'minimal':
new_height = self.constrain_to_multiple_of(scale_height * height)
new_width = self.constrain_to_multiple_of(scale_width * width)
else:
raise ValueError(
f'resize_method {self.__resize_method} not implemented')
return (new_width, new_height)
def __call__(self, sample):
width, height = self.get_size(sample['image'].shape[1],
sample['image'].shape[0])
# resize sample
sample['image'] = cv2.resize(
sample['image'],
(width, height),
interpolation=self.__image_interpolation_method,
)
if self.__resize_target:
if 'disparity' in sample:
sample['disparity'] = cv2.resize(
sample['disparity'],
(width, height),
interpolation=cv2.INTER_NEAREST,
)
if 'depth' in sample:
sample['depth'] = cv2.resize(sample['depth'], (width, height),
interpolation=cv2.INTER_NEAREST)
sample['mask'] = cv2.resize(
sample['mask'].astype(np.float32),
(width, height),
interpolation=cv2.INTER_NEAREST,
)
sample['mask'] = sample['mask'].astype(bool)
return sample
class NormalizeImage(object):
"""Normlize image by given mean and std.
"""
def __init__(self, mean, std):
self.__mean = mean
self.__std = std
def __call__(self, sample):
sample['image'] = (sample['image'] - self.__mean) / self.__std
return sample
class PrepareForNet(object):
"""Prepare sample for usage as network input.
"""
def __init__(self):
pass
def __call__(self, sample):
image = np.transpose(sample['image'], (2, 0, 1))
sample['image'] = np.ascontiguousarray(image).astype(np.float32)
if 'mask' in sample:
sample['mask'] = sample['mask'].astype(np.float32)
sample['mask'] = np.ascontiguousarray(sample['mask'])
if 'disparity' in sample:
disparity = sample['disparity'].astype(np.float32)
sample['disparity'] = np.ascontiguousarray(disparity)
if 'depth' in sample:
depth = sample['depth'].astype(np.float32)
sample['depth'] = np.ascontiguousarray(depth)
return sample

193
preprocessing/midas/utils.py Normal file
View File

@@ -0,0 +1,193 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
"""Utils for monoDepth."""
import re
import sys
import cv2
import numpy as np
import torch
def read_pfm(path):
"""Read pfm file.
Args:
path (str): path to file
Returns:
tuple: (data, scale)
"""
with open(path, 'rb') as file:
color = None
width = None
height = None
scale = None
endian = None
header = file.readline().rstrip()
if header.decode('ascii') == 'PF':
color = True
elif header.decode('ascii') == 'Pf':
color = False
else:
raise Exception('Not a PFM file: ' + path)
dim_match = re.match(r'^(\d+)\s(\d+)\s$',
file.readline().decode('ascii'))
if dim_match:
width, height = list(map(int, dim_match.groups()))
else:
raise Exception('Malformed PFM header.')
scale = float(file.readline().decode('ascii').rstrip())
if scale < 0:
# little-endian
endian = '<'
scale = -scale
else:
# big-endian
endian = '>'
data = np.fromfile(file, endian + 'f')
shape = (height, width, 3) if color else (height, width)
data = np.reshape(data, shape)
data = np.flipud(data)
return data, scale
def write_pfm(path, image, scale=1):
"""Write pfm file.
Args:
path (str): pathto file
image (array): data
scale (int, optional): Scale. Defaults to 1.
"""
with open(path, 'wb') as file:
color = None
if image.dtype.name != 'float32':
raise Exception('Image dtype must be float32.')
image = np.flipud(image)
if len(image.shape) == 3 and image.shape[2] == 3: # color image
color = True
elif (len(image.shape) == 2
or len(image.shape) == 3 and image.shape[2] == 1): # greyscale
color = False
else:
raise Exception(
'Image must have H x W x 3, H x W x 1 or H x W dimensions.')
file.write('PF\n' if color else 'Pf\n'.encode())
file.write('%d %d\n'.encode() % (image.shape[1], image.shape[0]))
endian = image.dtype.byteorder
if endian == '<' or endian == '=' and sys.byteorder == 'little':
scale = -scale
file.write('%f\n'.encode() % scale)
image.tofile(file)
def read_image(path):
"""Read image and output RGB image (0-1).
Args:
path (str): path to file
Returns:
array: RGB image (0-1)
"""
img = cv2.imread(path)
if img.ndim == 2:
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0
return img
def resize_image(img):
"""Resize image and make it fit for network.
Args:
img (array): image
Returns:
tensor: data ready for network
"""
height_orig = img.shape[0]
width_orig = img.shape[1]
if width_orig > height_orig:
scale = width_orig / 384
else:
scale = height_orig / 384
height = (np.ceil(height_orig / scale / 32) * 32).astype(int)
width = (np.ceil(width_orig / scale / 32) * 32).astype(int)
img_resized = cv2.resize(img, (width, height),
interpolation=cv2.INTER_AREA)
img_resized = (torch.from_numpy(np.transpose(
img_resized, (2, 0, 1))).contiguous().float())
img_resized = img_resized.unsqueeze(0)
return img_resized
def resize_depth(depth, width, height):
"""Resize depth map and bring to CPU (numpy).
Args:
depth (tensor): depth
width (int): image width
height (int): image height
Returns:
array: processed depth
"""
depth = torch.squeeze(depth[0, :, :, :]).to('cpu')
depth_resized = cv2.resize(depth.numpy(), (width, height),
interpolation=cv2.INTER_CUBIC)
return depth_resized
def write_depth(path, depth, bits=1):
"""Write depth map to pfm and png file.
Args:
path (str): filepath without extension
depth (array): depth
"""
write_pfm(path + '.pfm', depth.astype(np.float32))
depth_min = depth.min()
depth_max = depth.max()
max_val = (2**(8 * bits)) - 1
if depth_max - depth_min > np.finfo('float').eps:
out = max_val * (depth - depth_min) / (depth_max - depth_min)
else:
out = np.zeros(depth.shape, dtype=depth.type)
if bits == 1:
cv2.imwrite(path + '.png', out.astype('uint8'))
elif bits == 2:
cv2.imwrite(path + '.png', out.astype('uint16'))
return

510
preprocessing/midas/vit.py Normal file
View File

@@ -0,0 +1,510 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import math
import types
import timm
import torch
import torch.nn as nn
import torch.nn.functional as F
class Slice(nn.Module):
def __init__(self, start_index=1):
super(Slice, self).__init__()
self.start_index = start_index
def forward(self, x):
return x[:, self.start_index:]
class AddReadout(nn.Module):
def __init__(self, start_index=1):
super(AddReadout, self).__init__()
self.start_index = start_index
def forward(self, x):
if self.start_index == 2:
readout = (x[:, 0] + x[:, 1]) / 2
else:
readout = x[:, 0]
return x[:, self.start_index:] + readout.unsqueeze(1)
class ProjectReadout(nn.Module):
def __init__(self, in_features, start_index=1):
super(ProjectReadout, self).__init__()
self.start_index = start_index
self.project = nn.Sequential(nn.Linear(2 * in_features, in_features),
nn.GELU())
def forward(self, x):
readout = x[:, 0].unsqueeze(1).expand_as(x[:, self.start_index:])
features = torch.cat((x[:, self.start_index:], readout), -1)
return self.project(features)
class Transpose(nn.Module):
def __init__(self, dim0, dim1):
super(Transpose, self).__init__()
self.dim0 = dim0
self.dim1 = dim1
def forward(self, x):
x = x.transpose(self.dim0, self.dim1)
return x
def forward_vit(pretrained, x):
b, c, h, w = x.shape
_ = pretrained.model.forward_flex(x)
layer_1 = pretrained.activations['1']
layer_2 = pretrained.activations['2']
layer_3 = pretrained.activations['3']
layer_4 = pretrained.activations['4']
layer_1 = pretrained.act_postprocess1[0:2](layer_1)
layer_2 = pretrained.act_postprocess2[0:2](layer_2)
layer_3 = pretrained.act_postprocess3[0:2](layer_3)
layer_4 = pretrained.act_postprocess4[0:2](layer_4)
unflatten = nn.Sequential(
nn.Unflatten(
2,
torch.Size([
h // pretrained.model.patch_size[1],
w // pretrained.model.patch_size[0],
]),
))
if layer_1.ndim == 3:
layer_1 = unflatten(layer_1)
if layer_2.ndim == 3:
layer_2 = unflatten(layer_2)
if layer_3.ndim == 3:
layer_3 = unflatten(layer_3)
if layer_4.ndim == 3:
layer_4 = unflatten(layer_4)
layer_1 = pretrained.act_postprocess1[3:len(pretrained.act_postprocess1)](
layer_1)
layer_2 = pretrained.act_postprocess2[3:len(pretrained.act_postprocess2)](
layer_2)
layer_3 = pretrained.act_postprocess3[3:len(pretrained.act_postprocess3)](
layer_3)
layer_4 = pretrained.act_postprocess4[3:len(pretrained.act_postprocess4)](
layer_4)
return layer_1, layer_2, layer_3, layer_4
def _resize_pos_embed(self, posemb, gs_h, gs_w):
posemb_tok, posemb_grid = (
posemb[:, :self.start_index],
posemb[0, self.start_index:],
)
gs_old = int(math.sqrt(len(posemb_grid)))
posemb_grid = posemb_grid.reshape(1, gs_old, gs_old,
-1).permute(0, 3, 1, 2)
posemb_grid = F.interpolate(posemb_grid,
size=(gs_h, gs_w),
mode='bilinear')
posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_h * gs_w, -1)
posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
return posemb
def forward_flex(self, x):
b, c, h, w = x.shape
pos_embed = self._resize_pos_embed(self.pos_embed, h // self.patch_size[1],
w // self.patch_size[0])
B = x.shape[0]
if hasattr(self.patch_embed, 'backbone'):
x = self.patch_embed.backbone(x)
if isinstance(x, (list, tuple)):
x = x[
-1] # last feature if backbone outputs list/tuple of features
x = self.patch_embed.proj(x).flatten(2).transpose(1, 2)
if getattr(self, 'dist_token', None) is not None:
cls_tokens = self.cls_token.expand(
B, -1, -1) # stole cls_tokens impl from Phil Wang, thanks
dist_token = self.dist_token.expand(B, -1, -1)
x = torch.cat((cls_tokens, dist_token, x), dim=1)
else:
cls_tokens = self.cls_token.expand(
B, -1, -1) # stole cls_tokens impl from Phil Wang, thanks
x = torch.cat((cls_tokens, x), dim=1)
x = x + pos_embed
x = self.pos_drop(x)
for blk in self.blocks:
x = blk(x)
x = self.norm(x)
return x
activations = {}
def get_activation(name):
def hook(model, input, output):
activations[name] = output
return hook
def get_readout_oper(vit_features, features, use_readout, start_index=1):
if use_readout == 'ignore':
readout_oper = [Slice(start_index)] * len(features)
elif use_readout == 'add':
readout_oper = [AddReadout(start_index)] * len(features)
elif use_readout == 'project':
readout_oper = [
ProjectReadout(vit_features, start_index) for out_feat in features
]
else:
assert (
False
), "wrong operation for readout token, use_readout can be 'ignore', 'add', or 'project'"
return readout_oper
def _make_vit_b16_backbone(
model,
features=[96, 192, 384, 768],
size=[384, 384],
hooks=[2, 5, 8, 11],
vit_features=768,
use_readout='ignore',
start_index=1,
):
pretrained = nn.Module()
pretrained.model = model
pretrained.model.blocks[hooks[0]].register_forward_hook(
get_activation('1'))
pretrained.model.blocks[hooks[1]].register_forward_hook(
get_activation('2'))
pretrained.model.blocks[hooks[2]].register_forward_hook(
get_activation('3'))
pretrained.model.blocks[hooks[3]].register_forward_hook(
get_activation('4'))
pretrained.activations = activations
readout_oper = get_readout_oper(vit_features, features, use_readout,
start_index)
# 32, 48, 136, 384
pretrained.act_postprocess1 = nn.Sequential(
readout_oper[0],
Transpose(1, 2),
nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
nn.Conv2d(
in_channels=vit_features,
out_channels=features[0],
kernel_size=1,
stride=1,
padding=0,
),
nn.ConvTranspose2d(
in_channels=features[0],
out_channels=features[0],
kernel_size=4,
stride=4,
padding=0,
bias=True,
dilation=1,
groups=1,
),
)
pretrained.act_postprocess2 = nn.Sequential(
readout_oper[1],
Transpose(1, 2),
nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
nn.Conv2d(
in_channels=vit_features,
out_channels=features[1],
kernel_size=1,
stride=1,
padding=0,
),
nn.ConvTranspose2d(
in_channels=features[1],
out_channels=features[1],
kernel_size=2,
stride=2,
padding=0,
bias=True,
dilation=1,
groups=1,
),
)
pretrained.act_postprocess3 = nn.Sequential(
readout_oper[2],
Transpose(1, 2),
nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
nn.Conv2d(
in_channels=vit_features,
out_channels=features[2],
kernel_size=1,
stride=1,
padding=0,
),
)
pretrained.act_postprocess4 = nn.Sequential(
readout_oper[3],
Transpose(1, 2),
nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
nn.Conv2d(
in_channels=vit_features,
out_channels=features[3],
kernel_size=1,
stride=1,
padding=0,
),
nn.Conv2d(
in_channels=features[3],
out_channels=features[3],
kernel_size=3,
stride=2,
padding=1,
),
)
pretrained.model.start_index = start_index
pretrained.model.patch_size = [16, 16]
# We inject this function into the VisionTransformer instances so that
# we can use it with interpolated position embeddings without modifying the library source.
pretrained.model.forward_flex = types.MethodType(forward_flex,
pretrained.model)
pretrained.model._resize_pos_embed = types.MethodType(
_resize_pos_embed, pretrained.model)
return pretrained
def _make_pretrained_vitl16_384(pretrained, use_readout='ignore', hooks=None):
model = timm.create_model('vit_large_patch16_384', pretrained=pretrained)
hooks = [5, 11, 17, 23] if hooks is None else hooks
return _make_vit_b16_backbone(
model,
features=[256, 512, 1024, 1024],
hooks=hooks,
vit_features=1024,
use_readout=use_readout,
)
def _make_pretrained_vitb16_384(pretrained, use_readout='ignore', hooks=None):
model = timm.create_model('vit_base_patch16_384', pretrained=pretrained)
hooks = [2, 5, 8, 11] if hooks is None else hooks
return _make_vit_b16_backbone(model,
features=[96, 192, 384, 768],
hooks=hooks,
use_readout=use_readout)
def _make_pretrained_deitb16_384(pretrained, use_readout='ignore', hooks=None):
model = timm.create_model('vit_deit_base_patch16_384',
pretrained=pretrained)
hooks = [2, 5, 8, 11] if hooks is None else hooks
return _make_vit_b16_backbone(model,
features=[96, 192, 384, 768],
hooks=hooks,
use_readout=use_readout)
def _make_pretrained_deitb16_distil_384(pretrained,
use_readout='ignore',
hooks=None):
model = timm.create_model('vit_deit_base_distilled_patch16_384',
pretrained=pretrained)
hooks = [2, 5, 8, 11] if hooks is None else hooks
return _make_vit_b16_backbone(
model,
features=[96, 192, 384, 768],
hooks=hooks,
use_readout=use_readout,
start_index=2,
)
def _make_vit_b_rn50_backbone(
model,
features=[256, 512, 768, 768],
size=[384, 384],
hooks=[0, 1, 8, 11],
vit_features=768,
use_vit_only=False,
use_readout='ignore',
start_index=1,
):
pretrained = nn.Module()
pretrained.model = model
if use_vit_only is True:
pretrained.model.blocks[hooks[0]].register_forward_hook(
get_activation('1'))
pretrained.model.blocks[hooks[1]].register_forward_hook(
get_activation('2'))
else:
pretrained.model.patch_embed.backbone.stages[0].register_forward_hook(
get_activation('1'))
pretrained.model.patch_embed.backbone.stages[1].register_forward_hook(
get_activation('2'))
pretrained.model.blocks[hooks[2]].register_forward_hook(
get_activation('3'))
pretrained.model.blocks[hooks[3]].register_forward_hook(
get_activation('4'))
pretrained.activations = activations
readout_oper = get_readout_oper(vit_features, features, use_readout,
start_index)
if use_vit_only is True:
pretrained.act_postprocess1 = nn.Sequential(
readout_oper[0],
Transpose(1, 2),
nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
nn.Conv2d(
in_channels=vit_features,
out_channels=features[0],
kernel_size=1,
stride=1,
padding=0,
),
nn.ConvTranspose2d(
in_channels=features[0],
out_channels=features[0],
kernel_size=4,
stride=4,
padding=0,
bias=True,
dilation=1,
groups=1,
),
)
pretrained.act_postprocess2 = nn.Sequential(
readout_oper[1],
Transpose(1, 2),
nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
nn.Conv2d(
in_channels=vit_features,
out_channels=features[1],
kernel_size=1,
stride=1,
padding=0,
),
nn.ConvTranspose2d(
in_channels=features[1],
out_channels=features[1],
kernel_size=2,
stride=2,
padding=0,
bias=True,
dilation=1,
groups=1,
),
)
else:
pretrained.act_postprocess1 = nn.Sequential(nn.Identity(),
nn.Identity(),
nn.Identity())
pretrained.act_postprocess2 = nn.Sequential(nn.Identity(),
nn.Identity(),
nn.Identity())
pretrained.act_postprocess3 = nn.Sequential(
readout_oper[2],
Transpose(1, 2),
nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
nn.Conv2d(
in_channels=vit_features,
out_channels=features[2],
kernel_size=1,
stride=1,
padding=0,
),
)
pretrained.act_postprocess4 = nn.Sequential(
readout_oper[3],
Transpose(1, 2),
nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
nn.Conv2d(
in_channels=vit_features,
out_channels=features[3],
kernel_size=1,
stride=1,
padding=0,
),
nn.Conv2d(
in_channels=features[3],
out_channels=features[3],
kernel_size=3,
stride=2,
padding=1,
),
)
pretrained.model.start_index = start_index
pretrained.model.patch_size = [16, 16]
# We inject this function into the VisionTransformer instances so that
# we can use it with interpolated position embeddings without modifying the library source.
pretrained.model.forward_flex = types.MethodType(forward_flex,
pretrained.model)
# We inject this function into the VisionTransformer instances so that
# we can use it with interpolated position embeddings without modifying the library source.
pretrained.model._resize_pos_embed = types.MethodType(
_resize_pos_embed, pretrained.model)
return pretrained
def _make_pretrained_vitb_rn50_384(pretrained,
use_readout='ignore',
hooks=None,
use_vit_only=False):
model = timm.create_model('vit_base_resnet50_384', pretrained=pretrained)
hooks = [0, 1, 8, 11] if hooks is None else hooks
return _make_vit_b_rn50_backbone(
model,
features=[256, 512, 768, 768],
size=[384, 384],
hooks=hooks,
use_vit_only=use_vit_only,
use_readout=use_readout,
)

17
requirements.txt
View File

@@ -11,12 +11,21 @@ easydict
ftfy
dashscope
imageio-ffmpeg
# flash_attn
gradio>=5.0.0
# flash_attn
gradio>=5.0.0
numpy>=1.23.5,<2
einops
moviepy==1.0.3
mmgp==3.3.4
mmgp==3.4.0
peft==0.14.0
mutagen
pydantic==2.10.6
pydantic==2.10.6
decord
onnxruntime-gpu
rembg[gpu]==2.0.65
matplotlib
timm
segment-anything
omegaconf
hydra-core
# rembg==2.0.65

133
rife/IFNet_HDv3.py Normal file
View File

@@ -0,0 +1,133 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
# from ..model.warplayer import warp
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
backwarp_tenGrid = {}
def warp(tenInput, tenFlow, device):
k = (str(tenFlow.device), str(tenFlow.size()))
if k not in backwarp_tenGrid:
tenHorizontal = torch.linspace(-1.0, 1.0, tenFlow.shape[3], device=device).view(
1, 1, 1, tenFlow.shape[3]).expand(tenFlow.shape[0], -1, tenFlow.shape[2], -1)
tenVertical = torch.linspace(-1.0, 1.0, tenFlow.shape[2], device=device).view(
1, 1, tenFlow.shape[2], 1).expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3])
backwarp_tenGrid[k] = torch.cat(
[tenHorizontal, tenVertical], 1).to(device)
tenFlow = torch.cat([tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0),
tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0)], 1)
g = (backwarp_tenGrid[k] + tenFlow).permute(0, 2, 3, 1)
return torch.nn.functional.grid_sample(input=tenInput, grid=g, mode='bilinear', padding_mode='border', align_corners=True)
def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
return nn.Sequential(
nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
padding=padding, dilation=dilation, bias=True),
nn.PReLU(out_planes)
)
def conv_bn(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
return nn.Sequential(
nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
padding=padding, dilation=dilation, bias=False),
nn.BatchNorm2d(out_planes),
nn.PReLU(out_planes)
)
class IFBlock(nn.Module):
def __init__(self, in_planes, c=64):
super(IFBlock, self).__init__()
self.conv0 = nn.Sequential(
conv(in_planes, c//2, 3, 2, 1),
conv(c//2, c, 3, 2, 1),
)
self.convblock0 = nn.Sequential(
conv(c, c),
conv(c, c)
)
self.convblock1 = nn.Sequential(
conv(c, c),
conv(c, c)
)
self.convblock2 = nn.Sequential(
conv(c, c),
conv(c, c)
)
self.convblock3 = nn.Sequential(
conv(c, c),
conv(c, c)
)
self.conv1 = nn.Sequential(
nn.ConvTranspose2d(c, c//2, 4, 2, 1),
nn.PReLU(c//2),
nn.ConvTranspose2d(c//2, 4, 4, 2, 1),
)
self.conv2 = nn.Sequential(
nn.ConvTranspose2d(c, c//2, 4, 2, 1),
nn.PReLU(c//2),
nn.ConvTranspose2d(c//2, 1, 4, 2, 1),
)
def forward(self, x, flow, scale=1):
x = F.interpolate(x, scale_factor= 1. / scale, mode="bilinear", align_corners=False, recompute_scale_factor=False)
flow = F.interpolate(flow, scale_factor= 1. / scale, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 1. / scale
feat = self.conv0(torch.cat((x, flow), 1))
feat = self.convblock0(feat) + feat
feat = self.convblock1(feat) + feat
feat = self.convblock2(feat) + feat
feat = self.convblock3(feat) + feat
flow = self.conv1(feat)
mask = self.conv2(feat)
flow = F.interpolate(flow, scale_factor=scale, mode="bilinear", align_corners=False, recompute_scale_factor=False) * scale
mask = F.interpolate(mask, scale_factor=scale, mode="bilinear", align_corners=False, recompute_scale_factor=False)
return flow, mask
class IFNet(nn.Module):
def __init__(self):
super(IFNet, self).__init__()
self.block0 = IFBlock(7+4, c=90)
self.block1 = IFBlock(7+4, c=90)
self.block2 = IFBlock(7+4, c=90)
self.block_tea = IFBlock(10+4, c=90)
# self.contextnet = Contextnet()
# self.unet = Unet()
def forward(self, x, scale_list=[4, 2, 1], training=False):
if training == False:
channel = x.shape[1] // 2
img0 = x[:, :channel]
img1 = x[:, channel:]
flow_list = []
merged = []
mask_list = []
warped_img0 = img0
warped_img1 = img1
flow = (x[:, :4]).detach() * 0
mask = (x[:, :1]).detach() * 0
loss_cons = 0
block = [self.block0, self.block1, self.block2]
for i in range(3):
f0, m0 = block[i](torch.cat((warped_img0[:, :3], warped_img1[:, :3], mask), 1), flow, scale=scale_list[i])
f1, m1 = block[i](torch.cat((warped_img1[:, :3], warped_img0[:, :3], -mask), 1), torch.cat((flow[:, 2:4], flow[:, :2]), 1), scale=scale_list[i])
flow = flow + (f0 + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 2
mask = mask + (m0 + (-m1)) / 2
mask_list.append(mask)
flow_list.append(flow)
warped_img0 = warp(img0, flow[:, :2], device= flow.device)
warped_img1 = warp(img1, flow[:, 2:4], device= flow.device)
merged.append((warped_img0, warped_img1))
'''
c0 = self.contextnet(img0, flow[:, :2])
c1 = self.contextnet(img1, flow[:, 2:4])
tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)
res = tmp[:, 1:4] * 2 - 1
'''
for i in range(3):
mask_list[i] = torch.sigmoid(mask_list[i])
merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i])
# merged[i] = torch.clamp(merged[i] + res, 0, 1)
return flow_list, mask_list[2], merged

84
rife/RIFE_HDv3.py Normal file
View File

@@ -0,0 +1,84 @@
import torch
import torch.nn as nn
import numpy as np
from torch.optim import AdamW
import torch.optim as optim
import itertools
from torch.nn.parallel import DistributedDataParallel as DDP
from .IFNet_HDv3 import *
import torch.nn.functional as F
# from ..model.loss import *
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class Model:
def __init__(self, local_rank=-1):
self.flownet = IFNet()
# self.device()
# self.optimG = AdamW(self.flownet.parameters(), lr=1e-6, weight_decay=1e-4)
# self.epe = EPE()
# self.vgg = VGGPerceptualLoss().to(device)
# self.sobel = SOBEL()
if local_rank != -1:
self.flownet = DDP(self.flownet, device_ids=[local_rank], output_device=local_rank)
def train(self):
self.flownet.train()
def eval(self):
self.flownet.eval()
def to(self, device):
self.flownet.to(device)
def load_model(self, path, rank=0, device = "cuda"):
self.device = device
def convert(param):
if rank == -1:
return {
k.replace("module.", ""): v
for k, v in param.items()
if "module." in k
}
else:
return param
self.flownet.load_state_dict(convert(torch.load(path, map_location=device)))
def save_model(self, path, rank=0):
if rank == 0:
torch.save(self.flownet.state_dict(),'{}/flownet.pkl'.format(path))
def inference(self, img0, img1, scale=1.0):
imgs = torch.cat((img0, img1), 1)
scale_list = [4/scale, 2/scale, 1/scale]
flow, mask, merged = self.flownet(imgs, scale_list)
return merged[2]
def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None):
for param_group in self.optimG.param_groups:
param_group['lr'] = learning_rate
img0 = imgs[:, :3]
img1 = imgs[:, 3:]
if training:
self.train()
else:
self.eval()
scale = [4, 2, 1]
flow, mask, merged = self.flownet(torch.cat((imgs, gt), 1), scale=scale, training=training)
loss_l1 = (merged[2] - gt).abs().mean()
loss_smooth = self.sobel(flow[2], flow[2]*0).mean()
# loss_vgg = self.vgg(merged[2], gt)
if training:
self.optimG.zero_grad()
loss_G = loss_cons + loss_smooth * 0.1
loss_G.backward()
self.optimG.step()
else:
flow_teacher = flow[2]
return merged[2], {
'mask': mask,
'flow': flow[2][:, :2],
'loss_l1': loss_l1,
'loss_cons': loss_cons,
'loss_smooth': loss_smooth,
}

119
rife/inference.py Normal file
View File

@@ -0,0 +1,119 @@
import os
import torch
from torch.nn import functional as F
# from .model.pytorch_msssim import ssim_matlab
from .ssim import ssim_matlab
from .RIFE_HDv3 import Model
def get_frame(frames, frame_no):
if frame_no >= frames.shape[1]:
return None
frame = (frames[:, frame_no] + 1) /2
frame = frame.clip(0., 1.)
return frame
def add_frame(frames, frame, h, w):
frame = (frame * 2) - 1
frame = frame.clip(-1., 1.)
frame = frame.squeeze(0)
frame = frame[:, :h, :w]
frame = frame.unsqueeze(1)
frames.append(frame.cpu())
def process_frames(model, device, frames, exp):
pos = 0
output_frames = []
lastframe = get_frame(frames, 0)
_, h, w = lastframe.shape
scale = 1
fp16 = False
def make_inference(I0, I1, n):
middle = model.inference(I0, I1, scale)
if n == 1:
return [middle]
first_half = make_inference(I0, middle, n=n//2)
second_half = make_inference(middle, I1, n=n//2)
if n%2:
return [*first_half, middle, *second_half]
else:
return [*first_half, *second_half]
tmp = max(32, int(32 / scale))
ph = ((h - 1) // tmp + 1) * tmp
pw = ((w - 1) // tmp + 1) * tmp
padding = (0, pw - w, 0, ph - h)
def pad_image(img):
if(fp16):
return F.pad(img, padding).half()
else:
return F.pad(img, padding)
I1 = lastframe.to(device, non_blocking=True).unsqueeze(0)
I1 = pad_image(I1)
temp = None # save lastframe when processing static frame
while True:
if temp is not None:
frame = temp
temp = None
else:
pos += 1
frame = get_frame(frames, pos)
if frame is None:
break
I0 = I1
I1 = frame.to(device, non_blocking=True).unsqueeze(0)
I1 = pad_image(I1)
I0_small = F.interpolate(I0, (32, 32), mode='bilinear', align_corners=False)
I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False)
ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3])
break_flag = False
if ssim > 0.996 or pos > 100:
pos += 1
frame = get_frame(frames, pos)
if frame is None:
break_flag = True
frame = lastframe
else:
temp = frame
I1 = frame.to(device, non_blocking=True).unsqueeze(0)
I1 = pad_image(I1)
I1 = model.inference(I0, I1, scale)
I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False)
ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3])
frame = I1[0][:, :h, :w]
if ssim < 0.2:
output = []
for _ in range((2 ** exp) - 1):
output.append(I0)
else:
output = make_inference(I0, I1, 2**exp-1) if exp else []
add_frame(output_frames, lastframe, h, w)
for mid in output:
add_frame(output_frames, mid, h, w)
lastframe = frame
if break_flag:
break
add_frame(output_frames, lastframe, h, w)
return torch.cat( output_frames, dim=1)
def temporal_interpolation(model_path, frames, exp, device ="cuda"):
model = Model()
model.load_model(model_path, -1, device=device)
model.eval()
model.to(device=device)
with torch.no_grad():
output = process_frames(model, device, frames, exp)
return output

200
rife/ssim.py Normal file
View File

@@ -0,0 +1,200 @@
import torch
import torch.nn.functional as F
from math import exp
import numpy as np
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def gaussian(window_size, sigma):
gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
return gauss/gauss.sum()
def create_window(window_size, channel=1):
_1D_window = gaussian(window_size, 1.5).unsqueeze(1)
_2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0).to(device)
window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
return window
def create_window_3d(window_size, channel=1):
_1D_window = gaussian(window_size, 1.5).unsqueeze(1)
_2D_window = _1D_window.mm(_1D_window.t())
_3D_window = _2D_window.unsqueeze(2) @ (_1D_window.t())
window = _3D_window.expand(1, channel, window_size, window_size, window_size).contiguous().to(device)
return window
def ssim(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None):
# Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh).
if val_range is None:
if torch.max(img1) > 128:
max_val = 255
else:
max_val = 1
if torch.min(img1) < -0.5:
min_val = -1
else:
min_val = 0
L = max_val - min_val
else:
L = val_range
padd = 0
(_, channel, height, width) = img1.size()
if window is None:
real_size = min(window_size, height, width)
window = create_window(real_size, channel=channel).to(img1.device)
# mu1 = F.conv2d(img1, window, padding=padd, groups=channel)
# mu2 = F.conv2d(img2, window, padding=padd, groups=channel)
mu1 = F.conv2d(F.pad(img1, (5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=channel)
mu2 = F.conv2d(F.pad(img2, (5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=channel)
mu1_sq = mu1.pow(2)
mu2_sq = mu2.pow(2)
mu1_mu2 = mu1 * mu2
sigma1_sq = F.conv2d(F.pad(img1 * img1, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu1_sq
sigma2_sq = F.conv2d(F.pad(img2 * img2, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu2_sq
sigma12 = F.conv2d(F.pad(img1 * img2, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu1_mu2
C1 = (0.01 * L) ** 2
C2 = (0.03 * L) ** 2
v1 = 2.0 * sigma12 + C2
v2 = sigma1_sq + sigma2_sq + C2
cs = torch.mean(v1 / v2) # contrast sensitivity
ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2)
if size_average:
ret = ssim_map.mean()
else:
ret = ssim_map.mean(1).mean(1).mean(1)
if full:
return ret, cs
return ret
def ssim_matlab(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None):
# Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh).
if val_range is None:
if torch.max(img1) > 128:
max_val = 255
else:
max_val = 1
if torch.min(img1) < -0.5:
min_val = -1
else:
min_val = 0
L = max_val - min_val
else:
L = val_range
padd = 0
(_, _, height, width) = img1.size()
if window is None:
real_size = min(window_size, height, width)
window = create_window_3d(real_size, channel=1).to(img1.device)
# Channel is set to 1 since we consider color images as volumetric images
img1 = img1.unsqueeze(1)
img2 = img2.unsqueeze(1)
mu1 = F.conv3d(F.pad(img1, (5, 5, 5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=1)
mu2 = F.conv3d(F.pad(img2, (5, 5, 5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=1)
mu1_sq = mu1.pow(2)
mu2_sq = mu2.pow(2)
mu1_mu2 = mu1 * mu2
sigma1_sq = F.conv3d(F.pad(img1 * img1, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu1_sq
sigma2_sq = F.conv3d(F.pad(img2 * img2, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu2_sq
sigma12 = F.conv3d(F.pad(img1 * img2, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu1_mu2
C1 = (0.01 * L) ** 2
C2 = (0.03 * L) ** 2
v1 = 2.0 * sigma12 + C2
v2 = sigma1_sq + sigma2_sq + C2
cs = torch.mean(v1 / v2) # contrast sensitivity
ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2)
if size_average:
ret = ssim_map.mean()
else:
ret = ssim_map.mean(1).mean(1).mean(1)
if full:
return ret, cs
return ret
def msssim(img1, img2, window_size=11, size_average=True, val_range=None, normalize=False):
device = img1.device
weights = torch.FloatTensor([0.0448, 0.2856, 0.3001, 0.2363, 0.1333]).to(device)
levels = weights.size()[0]
mssim = []
mcs = []
for _ in range(levels):
sim, cs = ssim(img1, img2, window_size=window_size, size_average=size_average, full=True, val_range=val_range)
mssim.append(sim)
mcs.append(cs)
img1 = F.avg_pool2d(img1, (2, 2))
img2 = F.avg_pool2d(img2, (2, 2))
mssim = torch.stack(mssim)
mcs = torch.stack(mcs)
# Normalize (to avoid NaNs during training unstable models, not compliant with original definition)
if normalize:
mssim = (mssim + 1) / 2
mcs = (mcs + 1) / 2
pow1 = mcs ** weights
pow2 = mssim ** weights
# From Matlab implementation https://ece.uwaterloo.ca/~z70wang/research/iwssim/
output = torch.prod(pow1[:-1] * pow2[-1])
return output
# Classes to re-use window
class SSIM(torch.nn.Module):
def __init__(self, window_size=11, size_average=True, val_range=None):
super(SSIM, self).__init__()
self.window_size = window_size
self.size_average = size_average
self.val_range = val_range
# Assume 3 channel for SSIM
self.channel = 3
self.window = create_window(window_size, channel=self.channel)
def forward(self, img1, img2):
(_, channel, _, _) = img1.size()
if channel == self.channel and self.window.dtype == img1.dtype:
window = self.window
else:
window = create_window(self.window_size, channel).to(img1.device).type(img1.dtype)
self.window = window
self.channel = channel
_ssim = ssim(img1, img2, window=window, window_size=self.window_size, size_average=self.size_average)
dssim = (1 - _ssim) / 2
return dssim
class MSSSIM(torch.nn.Module):
def __init__(self, window_size=11, size_average=True, channel=3):
super(MSSSIM, self).__init__()
self.window_size = window_size
self.size_average = size_average
self.channel = channel
def forward(self, img1, img2):
return msssim(img1, img2, window_size=self.window_size, size_average=self.size_average)

14
wan/configs/__init__.py
View File

@@ -40,3 +40,17 @@ SUPPORTED_SIZES = {
'i2v-14B': ('720*1280', '1280*720', '480*832', '832*480'),
't2i-14B': tuple(SIZE_CONFIGS.keys()),
}
VACE_SIZE_CONFIGS = {
'480*832': (480, 832),
'832*480': (832, 480),
}
VACE_MAX_AREA_CONFIGS = {
'480*832': 480 * 832,
'832*480': 832 * 480,
}
VACE_SUPPORTED_SIZES = {
'vace-1.3B': ('480*832', '832*480'),
}

98
wan/image2video.py
View File

@@ -25,8 +25,7 @@ from .utils.fm_solvers import (FlowDPMSolverMultistepScheduler,
get_sampling_sigmas, retrieve_timesteps)
from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
from wan.modules.posemb_layers import get_rotary_pos_embed
from PIL import Image
from wan.utils.utils import resize_lanczos
def optimized_scale(positive_flat, negative_flat):
@@ -40,7 +39,8 @@ def optimized_scale(positive_flat, negative_flat):
st_star = dot_product / squared_norm
return st_star
class WanI2V:
@@ -48,7 +48,6 @@ class WanI2V:
self,
config,
checkpoint_dir,
device_id=0,
rank=0,
t5_fsdp=False,
dit_fsdp=False,
@@ -58,6 +57,8 @@ class WanI2V:
i2v720p= True,
model_filename ="",
text_encoder_filename="",
quantizeTransformer = False,
dtype = torch.bfloat16
):
r"""
Initializes the image-to-video generation model components.
@@ -82,23 +83,22 @@ class WanI2V:
Enable initializing Transformer Model on CPU. Only works without FSDP or USP.
init_on_cpu (`bool`, *optional*, defaults to True):
"""
self.device = torch.device(f"cuda:{device_id}")
self.device = torch.device(f"cuda")
self.config = config
self.rank = rank
self.use_usp = use_usp
self.t5_cpu = t5_cpu
self.dtype = dtype
self.num_train_timesteps = config.num_train_timesteps
self.param_dtype = config.param_dtype
shard_fn = partial(shard_model, device_id=device_id)
# shard_fn = partial(shard_model, device_id=device_id)
self.text_encoder = T5EncoderModel(
text_len=config.text_len,
dtype=config.t5_dtype,
device=torch.device('cpu'),
checkpoint_path=text_encoder_filename,
tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer),
shard_fn=shard_fn if t5_fsdp else None,
shard_fn=None,
)
self.vae_stride = config.vae_stride
@@ -117,34 +117,16 @@ class WanI2V:
logging.info(f"Creating WanModel from {model_filename}")
from mmgp import offload
self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel, writable_tensors= False) #forcedConfigPath= "ckpts/config2.json",
self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel,do_quantize= quantizeTransformer, writable_tensors= False)
if self.dtype == torch.float16 and not "fp16" in model_filename:
self.model.to(self.dtype)
# offload.save_model(self.model, "i2v_720p_fp16.safetensors",do_quantize=True)
if self.dtype == torch.float16:
self.vae.model.to(self.dtype)
# offload.save_model(self.model, "wan2.1_Fun_InP_1.3B_bf16_bis.safetensors")
self.model.eval().requires_grad_(False)
if t5_fsdp or dit_fsdp or use_usp:
init_on_cpu = False
if use_usp:
from xfuser.core.distributed import \
get_sequence_parallel_world_size
from .distributed.xdit_context_parallel import (usp_attn_forward,
usp_dit_forward)
for block in self.model.blocks:
block.self_attn.forward = types.MethodType(
usp_attn_forward, block.self_attn)
self.model.forward = types.MethodType(usp_dit_forward, self.model)
self.sp_size = get_sequence_parallel_world_size()
else:
self.sp_size = 1
# if dist.is_initialized():
# dist.barrier()
# if dit_fsdp:
# self.model = shard_fn(self.model)
# else:
# if not init_on_cpu:
# self.model.to(self.device)
self.sample_neg_prompt = config.sample_neg_prompt
@@ -208,16 +190,16 @@ class WanI2V:
- H: Frame height (from max_area)
- W: Frame width from max_area)
"""
img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(self.device)
img = TF.to_tensor(img)
lat_frames = int((frame_num - 1) // self.vae_stride[0] + 1)
any_end_frame = img2 !=None
if any_end_frame:
any_end_frame = True
img2 = TF.to_tensor(img2).sub_(0.5).div_(0.5).to(self.device)
img2 = TF.to_tensor(img2)
if add_frames_for_end_image:
frame_num +=1
lat_frames = int((frame_num - 2) // self.vae_stride[0] + 2)
h, w = img.shape[1:]
aspect_ratio = h / w
lat_h = round(
@@ -229,8 +211,16 @@ class WanI2V:
h = lat_h * self.vae_stride[1]
w = lat_w * self.vae_stride[2]
clip_image_size = self.clip.model.image_size
img_interpolated = resize_lanczos(img, h, w).sub_(0.5).div_(0.5).unsqueeze(0).transpose(0,1).to(self.device, self.dtype)
img = resize_lanczos(img, clip_image_size, clip_image_size)
img = img.sub_(0.5).div_(0.5).to(self.device, self.dtype)
if img2!= None:
img_interpolated2 = resize_lanczos(img2, h, w).sub_(0.5).div_(0.5).unsqueeze(0).transpose(0,1).to(self.device, self.dtype)
img2 = resize_lanczos(img2, clip_image_size, clip_image_size)
img2 = img2.sub_(0.5).div_(0.5).to(self.device, self.dtype)
max_seq_len = lat_frames * lat_h * lat_w // ( self.patch_size[1] * self.patch_size[2])
max_seq_len = int(math.ceil(max_seq_len / self.sp_size)) * self.sp_size
seed = seed if seed >= 0 else random.randint(0, sys.maxsize)
seed_g = torch.Generator(device=self.device)
@@ -267,27 +257,26 @@ class WanI2V:
context = [t.to(self.device) for t in context]
context_null = [t.to(self.device) for t in context_null]
context = [u.to(self.dtype) for u in context]
context_null = [u.to(self.dtype) for u in context_null]
clip_context = self.clip.visual([img[:, None, :, :]])
if offload_model:
self.clip.model.cpu()
from mmgp import offload
offload.last_offload_obj.unload_all()
if any_end_frame:
img_interpolated = torch.nn.functional.interpolate(img[None].cpu(), size=(h, w), mode='bicubic').transpose(0, 1).to(torch.bfloat16)
img2_interpolated = torch.nn.functional.interpolate(img2[None].cpu(), size=(h, w), mode='bicubic').transpose(0, 1).to(torch.bfloat16)
mean2 = 0
enc= torch.concat([
img_interpolated,
torch.full( (3, frame_num-2, h, w), mean2, device="cpu", dtype= torch.bfloat16),
img2_interpolated,
torch.full( (3, frame_num-2, h, w), mean2, device=self.device, dtype= self.dtype),
img_interpolated2,
], dim=1).to(self.device)
else:
enc= torch.concat([
torch.nn.functional.interpolate(
img[None].cpu(), size=(h, w), mode='bicubic').transpose(0, 1).to(torch.bfloat16),
torch.zeros(3, frame_num-1, h, w, device="cpu", dtype= torch.bfloat16)
img_interpolated,
torch.zeros(3, frame_num-1, h, w, device=self.device, dtype= self.dtype)
], dim=1).to(self.device)
lat_y = self.vae.encode([enc], VAE_tile_size, any_end_frame= any_end_frame and add_frames_for_end_image)[0]
@@ -333,7 +322,8 @@ class WanI2V:
'seq_len': max_seq_len,
'y': [y],
'freqs' : freqs,
'pipeline' : self
'pipeline' : self,
'callback' : callback
}
arg_null = {
@@ -342,7 +332,8 @@ class WanI2V:
'seq_len': max_seq_len,
'y': [y],
'freqs' : freqs,
'pipeline' : self
'pipeline' : self,
'callback' : callback
}
arg_both= {
@@ -352,7 +343,8 @@ class WanI2V:
'seq_len': max_seq_len,
'y': [y],
'freqs' : freqs,
'pipeline' : self
'pipeline' : self,
'callback' : callback
}
if offload_model:
@@ -363,7 +355,7 @@ class WanI2V:
# self.model.to(self.device)
if callback != None:
callback(-1, None)
callback(-1, True)
for i, t in enumerate(tqdm(timesteps)):
offload.set_step_no_for_lora(self.model, i)
@@ -437,10 +429,10 @@ class WanI2V:
del timestep
if callback is not None:
callback(i, latent)
callback(i, False)
x0 = [latent.to(self.device, dtype=torch.bfloat16)]
x0 = [latent.to(self.device, dtype=self.dtype)]
if offload_model:
self.model.cpu()
@@ -451,7 +443,7 @@ class WanI2V:
video = self.vae.decode(x0, VAE_tile_size, any_end_frame= any_end_frame and add_frames_for_end_image)[0]
if any_end_frame and add_frames_for_end_image:
# video[:, -1:] = img2_interpolated
# video[:, -1:] = img_interpolated2
video = video[:, :-1]
else:

19
wan/modules/attention.py
View File

@@ -5,6 +5,11 @@ from mmgp import offload
import torch.nn.functional as F
try:
from xformers.ops import memory_efficient_attention
except ImportError:
memory_efficient_attention = None
try:
import flash_attn_interface
FLASH_ATTN_3_AVAILABLE = True
@@ -123,13 +128,13 @@ def get_attention_modes():
ret = ["sdpa", "auto"]
if flash_attn != None:
ret.append("flash")
# if memory_efficient_attention != None:
# ret.append("xformers")
if memory_efficient_attention != None:
ret.append("xformers")
if sageattn_varlen_wrapper != None:
ret.append("sage")
if sageattn != None and version("sageattention").startswith("2") :
ret.append("sage2")
return ret
def get_supported_attention_modes():
@@ -338,6 +343,14 @@ def pay_attention(
deterministic=deterministic).unflatten(0, (b, lq))
# output
elif attn=="xformers":
x = memory_efficient_attention(
q.unsqueeze(0),
k.unsqueeze(0),
v.unsqueeze(0),
) #.unsqueeze(0)
return x.type(out_dtype)

239
wan/modules/model.py
View File

@@ -77,73 +77,6 @@ def rope_params_riflex(max_seq_len, dim, theta=10000, L_test=30, k=6):
def rope_apply_(x, grid_sizes, freqs):
assert x.shape[0]==1
n, c = x.size(2), x.size(3) // 2
# split freqs
freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
f, h, w = grid_sizes[0]
seq_len = f * h * w
x_i = x[0, :seq_len, :, :]
x_i = x_i.to(torch.float32)
x_i = x_i.reshape(seq_len, n, -1, 2)
x_i = torch.view_as_complex(x_i)
freqs_i = torch.cat([
freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
], dim=-1)
freqs_i= freqs_i.reshape(seq_len, 1, -1)
# apply rotary embedding
x_i *= freqs_i
x_i = torch.view_as_real(x_i).flatten(2)
x[0, :seq_len, :, :] = x_i.to(torch.bfloat16)
# x_i = torch.cat([x_i, x[0, seq_len:]])
return x
# @amp.autocast(enabled=False)
def rope_apply(x, grid_sizes, freqs):
n, c = x.size(2), x.size(3) // 2
# split freqs
freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
# loop over samples
output = []
for i, (f, h, w) in enumerate(grid_sizes):
seq_len = f * h * w
# precompute multipliers
# x_i = x[i, :seq_len]
x_i = x[i]
x_i = x_i[:seq_len, :, :]
x_i = x_i.to(torch.float32)
x_i = x_i.reshape(seq_len, n, -1, 2)
x_i = torch.view_as_complex(x_i)
freqs_i = torch.cat([
freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
],
dim=-1).reshape(seq_len, 1, -1)
# apply rotary embedding
x_i *= freqs_i
x_i = torch.view_as_real(x_i).flatten(2)
x_i = x_i.to(torch.bfloat16)
x_i = torch.cat([x_i, x[i, seq_len:]])
# append to collection
output.append(x_i)
return torch.stack(output) #.float()
def relative_l1_distance(last_tensor, current_tensor):
l1_distance = torch.abs(last_tensor - current_tensor).mean()
norm = torch.abs(last_tensor).mean()
@@ -256,8 +189,6 @@ class WanSelfAttention(nn.Module):
k = k.view(b, s, n, d)
v = self.v(x).view(b, s, n, d)
del x
# rope_apply_(q, grid_sizes, freqs)
# rope_apply_(k, grid_sizes, freqs)
qklist = [q,k]
del q,k
q,k = apply_rotary_emb(qklist, freqs, head_first=False)
@@ -377,6 +308,7 @@ class WanI2VCrossAttention(WanSelfAttention):
return x
WAN_CROSSATTENTION_CLASSES = {
't2v_cross_attn': WanT2VCrossAttention,
'i2v_cross_attn': WanI2VCrossAttention,
@@ -393,7 +325,9 @@ class WanAttentionBlock(nn.Module):
window_size=(-1, -1),
qk_norm=True,
cross_attn_norm=False,
eps=1e-6):
eps=1e-6,
block_id=None
):
super().__init__()
self.dim = dim
self.ffn_dim = ffn_dim
@@ -422,6 +356,7 @@ class WanAttentionBlock(nn.Module):
# modulation
self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
self.block_id = block_id
def forward(
self,
@@ -432,6 +367,8 @@ class WanAttentionBlock(nn.Module):
freqs,
context,
context_lens,
hints= None,
context_scale=1.0,
):
r"""
Args:
@@ -441,6 +378,21 @@ class WanAttentionBlock(nn.Module):
grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
"""
hint = None
if self.block_id is not None and hints is not None:
kwargs = {
"seq_lens" : seq_lens,
"grid_sizes" : grid_sizes,
"freqs" :freqs,
"context" : context,
"context_lens" : context_lens,
"e" : e,
}
if self.block_id == 0:
hint = self.vace(hints, x, **kwargs)
else:
hint = self.vace(hints, None, **kwargs)
e = (self.modulation + e).chunk(6, dim=1)
# self-attention
@@ -461,7 +413,6 @@ class WanAttentionBlock(nn.Module):
y *= 1 + e[4]
y += e[3]
ffn = self.ffn[0]
gelu = self.ffn[1]
ffn2= self.ffn[2]
@@ -479,11 +430,51 @@ class WanAttentionBlock(nn.Module):
x.addcmul_(y, e[5])
return x
if hint is not None:
if context_scale == 1:
x.add_(hint)
else:
x.add_(hint, alpha= context_scale)
return x
class VaceWanAttentionBlock(WanAttentionBlock):
def __init__(
self,
cross_attn_type,
dim,
ffn_dim,
num_heads,
window_size=(-1, -1),
qk_norm=True,
cross_attn_norm=False,
eps=1e-6,
block_id=0
):
super().__init__(cross_attn_type, dim, ffn_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps)
self.block_id = block_id
if block_id == 0:
self.before_proj = nn.Linear(self.dim, self.dim)
nn.init.zeros_(self.before_proj.weight)
nn.init.zeros_(self.before_proj.bias)
self.after_proj = nn.Linear(self.dim, self.dim)
nn.init.zeros_(self.after_proj.weight)
nn.init.zeros_(self.after_proj.bias)
def forward(self, hints, x, **kwargs):
# behold dbm magic !
c = hints[0]
hints[0] = None
if self.block_id == 0:
c = self.before_proj(c)
c += x
c = super().forward(c, **kwargs)
c_skip = self.after_proj(c)
hints[0] = c
return c_skip
class Head(nn.Module):
def __init__(self, dim, out_dim, patch_size, eps=1e-6):
@@ -508,9 +499,9 @@ class Head(nn.Module):
e(Tensor): Shape [B, C]
"""
# assert e.dtype == torch.float32
dtype = x.dtype
e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)
x = self.norm(x).to(torch.bfloat16)
x = self.norm(x).to(dtype)
x *= (1 + e[1])
x += e[0]
x = self.head(x)
@@ -544,6 +535,8 @@ class WanModel(ModelMixin, ConfigMixin):
@register_to_config
def __init__(self,
vace_layers=None,
vace_in_dim=None,
model_type='t2v',
patch_size=(1, 2, 2),
text_len=512,
@@ -628,12 +621,13 @@ class WanModel(ModelMixin, ConfigMixin):
self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
# blocks
cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn'
self.blocks = nn.ModuleList([
WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,
window_size, qk_norm, cross_attn_norm, eps)
for _ in range(num_layers)
])
if vace_layers == None:
cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn'
self.blocks = nn.ModuleList([
WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,
window_size, qk_norm, cross_attn_norm, eps)
for _ in range(num_layers)
])
# head
self.head = Head(dim, out_dim, patch_size, eps)
@@ -646,6 +640,33 @@ class WanModel(ModelMixin, ConfigMixin):
# initialize weights
self.init_weights()
if vace_layers != None:
self.vace_layers = [i for i in range(0, self.num_layers, 2)] if vace_layers is None else vace_layers
self.vace_in_dim = self.in_dim if vace_in_dim is None else vace_in_dim
assert 0 in self.vace_layers
self.vace_layers_mapping = {i: n for n, i in enumerate(self.vace_layers)}
# blocks
self.blocks = nn.ModuleList([
WanAttentionBlock('t2v_cross_attn', self.dim, self.ffn_dim, self.num_heads, self.window_size, self.qk_norm,
self.cross_attn_norm, self.eps,
block_id=self.vace_layers_mapping[i] if i in self.vace_layers else None)
for i in range(self.num_layers)
])
# vace blocks
self.vace_blocks = nn.ModuleList([
VaceWanAttentionBlock('t2v_cross_attn', self.dim, self.ffn_dim, self.num_heads, self.window_size, self.qk_norm,
self.cross_attn_norm, self.eps, block_id=i)
for i in self.vace_layers
])
# vace patch embeddings
self.vace_patch_embedding = nn.Conv3d(
self.vace_in_dim, self.dim, kernel_size=self.patch_size, stride=self.patch_size
)
def compute_teacache_threshold(self, start_step, timesteps = None, speed_factor =0):
rescale_func = np.poly1d(self.coefficients)
@@ -688,6 +709,7 @@ class WanModel(ModelMixin, ConfigMixin):
self.rel_l1_thresh = best_threshold
print(f"Tea Cache, best threshold found:{best_threshold:0.2f} with gain x{len(timesteps)/(target_nb_steps - best_signed_diff):0.2f} for a target of x{speed_factor}")
return best_threshold
def forward(
self,
@@ -695,6 +717,8 @@ class WanModel(ModelMixin, ConfigMixin):
t,
context,
seq_len,
vace_context = None,
vace_context_scale=1.0,
clip_fea=None,
y=None,
freqs = None,
@@ -704,6 +728,7 @@ class WanModel(ModelMixin, ConfigMixin):
is_uncond=False,
max_steps = 0,
slg_layers=None,
callback = None,
):
r"""
Forward pass through the diffusion model
@@ -763,7 +788,7 @@ class WanModel(ModelMixin, ConfigMixin):
# time embeddings
e = self.time_embedding(
sinusoidal_embedding_1d(self.freq_dim, t))
e0 = self.time_projection(e).unflatten(1, (6, self.dim)).to(torch.bfloat16)
e0 = self.time_projection(e).unflatten(1, (6, self.dim)).to(e.dtype)
# context
context_lens = None
@@ -796,6 +821,34 @@ class WanModel(ModelMixin, ConfigMixin):
x_list = [x]
context_list = [context]
del x
# arguments
kwargs = dict(
seq_lens=seq_lens,
grid_sizes=grid_sizes,
freqs=freqs,
context_lens=context_lens,
)
if vace_context == None:
hints_list = [None ] *len(x_list)
else:
# embeddings
c = [self.vace_patch_embedding(u.unsqueeze(0)) for u in vace_context]
c = [u.flatten(2).transpose(1, 2) for u in c]
if (len(c) == 1 and seq_len == c[0].size(1)):
c = c[0]
else:
c = torch.cat([
torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],
dim=1) for u in c
])
kwargs['context_scale'] = vace_context_scale
hints_list = [ [c] for _ in range(len(x_list)) ]
del c
should_calc = True
if self.enable_teacache:
if is_uncond:
@@ -827,20 +880,11 @@ class WanModel(ModelMixin, ConfigMixin):
if joint_pass or not is_uncond:
self.previous_residual_cond = None
ori_hidden_states = x_list[0].clone()
# arguments
kwargs = dict(
# e=e0,
seq_lens=seq_lens,
grid_sizes=grid_sizes,
freqs=freqs,
# context=context,
context_lens=context_lens)
for block_idx, block in enumerate(self.blocks):
offload.shared_state["layer"] = block_idx
if "refresh" in offload.shared_state:
del offload.shared_state["refresh"]
offload.shared_state["callback"](-1, -1, True)
if callback != None:
callback(-1, False, True)
if pipeline._interrupt:
if joint_pass:
return None, None
@@ -853,9 +897,10 @@ class WanModel(ModelMixin, ConfigMixin):
x_list[0] = block(x_list[0], context = context_list[0], e= e0, **kwargs)
else:
for i, (x, context) in enumerate(zip(x_list, context_list)):
x_list[i] = block(x, context = context, e= e0, **kwargs)
for i, (x, context, hints) in enumerate(zip(x_list, context_list, hints_list)):
x_list[i] = block(x, context = context, hints= hints, e= e0, **kwargs)
del x
del context, hints
if self.enable_teacache:
if joint_pass:

11
wan/modules/vae.py
View File

@@ -51,10 +51,11 @@ class RMS_norm(nn.Module):
self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.
def forward(self, x):
dtype = x.dtype
x = F.normalize(
x, dim=(1 if self.channel_first else
-1)) * self.scale * self.gamma + self.bias
x = x.to(torch.bfloat16)
x = x.to(dtype)
return x
class Upsample(nn.Upsample):
@@ -208,6 +209,7 @@ class ResidualBlock(nn.Module):
def forward(self, x, feat_cache=None, feat_idx=[0]):
h = self.shortcut(x)
dtype = x.dtype
for layer in self.residual:
if isinstance(layer, CausalConv3d) and feat_cache is not None:
idx = feat_idx[0]
@@ -219,11 +221,11 @@ class ResidualBlock(nn.Module):
cache_x.device), cache_x
],
dim=2)
x = layer(x, feat_cache[idx]).to(torch.bfloat16)
x = layer(x, feat_cache[idx]).to(dtype)
feat_cache[idx] = cache_x#.to("cpu")
feat_idx[0] += 1
else:
x = layer(x).to(torch.bfloat16)
x = layer(x).to(dtype)
return x + h
@@ -323,6 +325,7 @@ class Encoder3d(nn.Module):
CausalConv3d(out_dim, z_dim, 3, padding=1))
def forward(self, x, feat_cache=None, feat_idx=[0]):
dtype = x.dtype
if feat_cache is not None:
idx = feat_idx[0]
cache_x = x[:, :, -CACHE_T:, :, :].clone()
@@ -333,7 +336,7 @@ class Encoder3d(nn.Module):
cache_x.device), cache_x
],
dim=2)
x = self.conv1(x, feat_cache[idx]).to(torch.bfloat16)
x = self.conv1(x, feat_cache[idx]).to(dtype)
feat_cache[idx] = cache_x
del cache_x
feat_idx[0] += 1

276
wan/text2video.py
View File

@@ -13,7 +13,9 @@ import torch
import torch.cuda.amp as amp
import torch.distributed as dist
from tqdm import tqdm
from PIL import Image
import torchvision.transforms.functional as TF
import torch.nn.functional as F
from .distributed.fsdp import shard_model
from .modules.model import WanModel
from .modules.t5 import T5EncoderModel
@@ -22,6 +24,7 @@ from .utils.fm_solvers import (FlowDPMSolverMultistepScheduler,
get_sampling_sigmas, retrieve_timesteps)
from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
from wan.modules.posemb_layers import get_rotary_pos_embed
from .utils.vace_preprocessor import VaceVideoProcessor
def optimized_scale(positive_flat, negative_flat):
@@ -44,14 +47,15 @@ class WanT2V:
self,
config,
checkpoint_dir,
device_id=0,
rank=0,
t5_fsdp=False,
dit_fsdp=False,
use_usp=False,
t5_cpu=False,
model_filename = None,
text_encoder_filename = None
text_encoder_filename = None,
quantizeTransformer = False,
dtype = torch.bfloat16
):
r"""
Initializes the Wan text-to-video generation model components.
@@ -74,25 +78,24 @@ class WanT2V:
t5_cpu (`bool`, *optional*, defaults to False):
Whether to place T5 model on CPU. Only works without t5_fsdp.
"""
self.device = torch.device(f"cuda:{device_id}")
self.device = torch.device(f"cuda")
self.config = config
self.rank = rank
self.t5_cpu = t5_cpu
self.dtype = dtype
self.num_train_timesteps = config.num_train_timesteps
self.param_dtype = config.param_dtype
shard_fn = partial(shard_model, device_id=device_id)
self.text_encoder = T5EncoderModel(
text_len=config.text_len,
dtype=config.t5_dtype,
device=torch.device('cpu'),
checkpoint_path=text_encoder_filename,
tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer),
shard_fn=shard_fn if t5_fsdp else None)
shard_fn= None)
self.vae_stride = config.vae_stride
self.patch_size = config.patch_size
self.patch_size = config.patch_size
self.vae = WanVAE(
@@ -102,38 +105,181 @@ class WanT2V:
logging.info(f"Creating WanModel from {model_filename}")
from mmgp import offload
self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel, writable_tensors= False)
self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel,do_quantize= quantizeTransformer, writable_tensors= False)
if self.dtype == torch.float16 and not "fp16" in model_filename:
self.model.to(self.dtype)
# offload.save_model(self.model, "t2v_fp16.safetensors",do_quantize=True)
if self.dtype == torch.float16:
self.vae.model.to(self.dtype)
self.model.eval().requires_grad_(False)
if use_usp:
from xfuser.core.distributed import \
get_sequence_parallel_world_size
from .distributed.xdit_context_parallel import (usp_attn_forward,
usp_dit_forward)
for block in self.model.blocks:
block.self_attn.forward = types.MethodType(
usp_attn_forward, block.self_attn)
self.model.forward = types.MethodType(usp_dit_forward, self.model)
self.sp_size = get_sequence_parallel_world_size()
else:
self.sp_size = 1
# if dist.is_initialized():
# dist.barrier()
# if dit_fsdp:
# self.model = shard_fn(self.model)
# else:
# self.model.to(self.device)
self.sample_neg_prompt = config.sample_neg_prompt
if "Vace" in model_filename:
self.vid_proc = VaceVideoProcessor(downsample=tuple([x * y for x, y in zip(config.vae_stride, self.patch_size)]),
min_area=480*832,
max_area=480*832,
min_fps=config.sample_fps,
max_fps=config.sample_fps,
zero_start=True,
seq_len=32760,
keep_last=True)
self.adapt_vace_model()
def vace_encode_frames(self, frames, ref_images, masks=None, tile_size = 0):
if ref_images is None:
ref_images = [None] * len(frames)
else:
assert len(frames) == len(ref_images)
if masks is None:
latents = self.vae.encode(frames, tile_size = tile_size)
else:
inactive = [i * (1 - m) + 0 * m for i, m in zip(frames, masks)]
reactive = [i * m + 0 * (1 - m) for i, m in zip(frames, masks)]
inactive = self.vae.encode(inactive, tile_size = tile_size)
reactive = self.vae.encode(reactive, tile_size = tile_size)
latents = [torch.cat((u, c), dim=0) for u, c in zip(inactive, reactive)]
cat_latents = []
for latent, refs in zip(latents, ref_images):
if refs is not None:
if masks is None:
ref_latent = self.vae.encode(refs, tile_size = tile_size)
else:
ref_latent = self.vae.encode(refs, tile_size = tile_size)
ref_latent = [torch.cat((u, torch.zeros_like(u)), dim=0) for u in ref_latent]
assert all([x.shape[1] == 1 for x in ref_latent])
latent = torch.cat([*ref_latent, latent], dim=1)
cat_latents.append(latent)
return cat_latents
def vace_encode_masks(self, masks, ref_images=None):
if ref_images is None:
ref_images = [None] * len(masks)
else:
assert len(masks) == len(ref_images)
result_masks = []
for mask, refs in zip(masks, ref_images):
c, depth, height, width = mask.shape
new_depth = int((depth + 3) // self.vae_stride[0])
height = 2 * (int(height) // (self.vae_stride[1] * 2))
width = 2 * (int(width) // (self.vae_stride[2] * 2))
# reshape
mask = mask[0, :, :, :]
mask = mask.view(
depth, height, self.vae_stride[1], width, self.vae_stride[1]
) # depth, height, 8, width, 8
mask = mask.permute(2, 4, 0, 1, 3) # 8, 8, depth, height, width
mask = mask.reshape(
self.vae_stride[1] * self.vae_stride[2], depth, height, width
) # 8*8, depth, height, width
# interpolation
mask = F.interpolate(mask.unsqueeze(0), size=(new_depth, height, width), mode='nearest-exact').squeeze(0)
if refs is not None:
length = len(refs)
mask_pad = torch.zeros_like(mask[:, :length, :, :])
mask = torch.cat((mask_pad, mask), dim=1)
result_masks.append(mask)
return result_masks
def vace_latent(self, z, m):
return [torch.cat([zz, mm], dim=0) for zz, mm in zip(z, m)]
def prepare_source(self, src_video, src_mask, src_ref_images, total_frames, image_size, device, original_video = False, keep_frames= [], start_frame = 0, pre_src_video = None):
image_sizes = []
trim_video = len(keep_frames)
for i, (sub_src_video, sub_src_mask, sub_pre_src_video) in enumerate(zip(src_video, src_mask,pre_src_video)):
prepend_count = 0 if sub_pre_src_video == None else sub_pre_src_video.shape[1]
num_frames = total_frames - prepend_count
if sub_src_mask is not None and sub_src_video is not None:
src_video[i], src_mask[i], _, _, _ = self.vid_proc.load_video_pair(sub_src_video, sub_src_mask, max_frames= num_frames, trim_video = trim_video - prepend_count, start_frame = start_frame)
# src_video is [-1, 1], 0 = inpainting area (in fact 127 in [0, 255])
# src_mask is [-1, 1], 0 = preserve original video (in fact 127 in [0, 255]) and 1 = Inpainting (in fact 255 in [0, 255])
src_video[i] = src_video[i].to(device)
src_mask[i] = src_mask[i].to(device)
if prepend_count > 0:
src_video[i] = torch.cat( [sub_pre_src_video, src_video[i]], dim=1)
src_mask[i] = torch.cat( [torch.zeros_like(sub_pre_src_video), src_mask[i]] ,1)
src_video_shape = src_video[i].shape
if src_video_shape[1] != total_frames:
src_video[i] = torch.cat( [src_video[i], src_video[i].new_zeros(src_video_shape[0], total_frames -src_video_shape[1], *src_video_shape[-2:])], dim=1)
src_mask[i] = torch.cat( [src_mask[i], src_mask[i].new_ones(src_video_shape[0], total_frames -src_video_shape[1], *src_video_shape[-2:])], dim=1)
src_mask[i] = torch.clamp((src_mask[i][:1, :, :, :] + 1) / 2, min=0, max=1)
image_sizes.append(src_video[i].shape[2:])
elif sub_src_video is None:
if prepend_count > 0:
src_video[i] = torch.cat( [sub_pre_src_video, torch.zeros((3, num_frames, image_size[0], image_size[1]), device=device)], dim=1)
src_mask[i] = torch.cat( [torch.zeros_like(sub_pre_src_video), torch.ones((3, num_frames, image_size[0], image_size[1]), device=device)] ,1)
else:
src_video[i] = torch.zeros((3, num_frames, image_size[0], image_size[1]), device=device)
src_mask[i] = torch.ones_like(src_video[i], device=device)
image_sizes.append(image_size)
else:
src_video[i], _, _, _ = self.vid_proc.load_video(sub_src_video, max_frames= num_frames, trim_video = trim_video - prepend_count, start_frame = start_frame)
src_video[i] = src_video[i].to(device)
src_mask[i] = torch.zeros_like(src_video[i], device=device) if original_video else torch.ones_like(src_video[i], device=device)
if prepend_count > 0:
src_video[i] = torch.cat( [sub_pre_src_video, src_video[i]], dim=1)
src_mask[i] = torch.cat( [torch.zeros_like(sub_pre_src_video), src_mask[i]] ,1)
src_video_shape = src_video[i].shape
if src_video_shape[1] != total_frames:
src_video[i] = torch.cat( [src_video[i], src_video[i].new_zeros(src_video_shape[0], total_frames -src_video_shape[1], *src_video_shape[-2:])], dim=1)
src_mask[i] = torch.cat( [src_mask[i], src_mask[i].new_ones(src_video_shape[0], total_frames -src_video_shape[1], *src_video_shape[-2:])], dim=1)
image_sizes.append(src_video[i].shape[2:])
for k, keep in enumerate(keep_frames):
if not keep:
src_video[i][:, k:k+1] = 0
src_mask[i][:, k:k+1] = 1
for i, ref_images in enumerate(src_ref_images):
if ref_images is not None:
image_size = image_sizes[i]
for j, ref_img in enumerate(ref_images):
if ref_img is not None:
ref_img = TF.to_tensor(ref_img).sub_(0.5).div_(0.5).unsqueeze(1)
if ref_img.shape[-2:] != image_size:
canvas_height, canvas_width = image_size
ref_height, ref_width = ref_img.shape[-2:]
white_canvas = torch.ones((3, 1, canvas_height, canvas_width), device=device) # [-1, 1]
scale = min(canvas_height / ref_height, canvas_width / ref_width)
new_height = int(ref_height * scale)
new_width = int(ref_width * scale)
resized_image = F.interpolate(ref_img.squeeze(1).unsqueeze(0), size=(new_height, new_width), mode='bilinear', align_corners=False).squeeze(0).unsqueeze(1)
top = (canvas_height - new_height) // 2
left = (canvas_width - new_width) // 2
white_canvas[:, :, top:top + new_height, left:left + new_width] = resized_image
ref_img = white_canvas
src_ref_images[i][j] = ref_img.to(device)
return src_video, src_mask, src_ref_images
def decode_latent(self, zs, ref_images=None, tile_size= 0 ):
if ref_images is None:
ref_images = [None] * len(zs)
else:
assert len(zs) == len(ref_images)
trimed_zs = []
for z, refs in zip(zs, ref_images):
if refs is not None:
z = z[:, len(refs):, :, :]
trimed_zs.append(z)
return self.vae.decode(trimed_zs, tile_size= tile_size)
def generate(self,
input_prompt,
input_frames= None,
input_masks = None,
input_ref_images = None,
context_scale=1.0,
size=(1280, 720),
frame_num=81,
shift=5.0,
@@ -187,14 +333,6 @@ class WanT2V:
- W: Frame width from size)
"""
# preprocess
F = frame_num
target_shape = (self.vae.model.z_dim, (F - 1) // self.vae_stride[0] + 1,
size[1] // self.vae_stride[1],
size[0] // self.vae_stride[2])
seq_len = math.ceil((target_shape[2] * target_shape[3]) /
(self.patch_size[1] * self.patch_size[2]) *
target_shape[1] / self.sp_size) * self.sp_size
if n_prompt == "":
n_prompt = self.sample_neg_prompt
@@ -213,6 +351,31 @@ class WanT2V:
context_null = self.text_encoder([n_prompt], torch.device('cpu'))
context = [t.to(self.device) for t in context]
context_null = [t.to(self.device) for t in context_null]
if input_frames != None:
# vace context encode
input_frames = [u.to(self.device) for u in input_frames]
input_ref_images = [ None if u == None else [v.to(self.device) for v in u] for u in input_ref_images]
input_masks = [u.to(self.device) for u in input_masks]
z0 = self.vace_encode_frames(input_frames, input_ref_images, masks=input_masks, tile_size = VAE_tile_size)
m0 = self.vace_encode_masks(input_masks, input_ref_images)
z = self.vace_latent(z0, m0)
target_shape = list(z0[0].shape)
target_shape[0] = int(target_shape[0] / 2)
else:
F = frame_num
target_shape = (self.vae.model.z_dim, (F - 1) // self.vae_stride[0] + 1,
size[1] // self.vae_stride[1],
size[0] // self.vae_stride[2])
seq_len = math.ceil((target_shape[2] * target_shape[3]) /
(self.patch_size[1] * self.patch_size[2]) *
target_shape[1])
context = [u.to(self.dtype) for u in context]
context_null = [u.to(self.dtype) for u in context_null]
noise = [
torch.randn(
@@ -261,14 +424,16 @@ class WanT2V:
arg_c = {'context': context, 'seq_len': seq_len, 'freqs': freqs, 'pipeline': self}
arg_null = {'context': context_null, 'seq_len': seq_len, 'freqs': freqs, 'pipeline': self}
arg_both = {'context': context, 'context2': context_null, 'seq_len': seq_len, 'freqs': freqs, 'pipeline': self}
if input_frames != None:
vace_dict = {'vace_context' : z, 'vace_context_scale' : context_scale}
arg_c.update(vace_dict)
arg_null.update(vace_dict)
arg_both.update(vace_dict)
# arg_c = {'context': context, 'seq_len': seq_len, 'freqs': freqs, 'pipeline': self, "max_steps": sampling_steps}
# arg_null = {'context': context_null, 'seq_len': seq_len, 'freqs': freqs, 'pipeline': self, "max_steps": sampling_steps}
# arg_both = {'context': context, 'context2': context_null, 'seq_len': seq_len, 'freqs': freqs, 'pipeline': self, "max_steps": sampling_steps}
if self.model.enable_teacache:
self.model.compute_teacache_threshold(self.model.teacache_start_step, timesteps, self.model.teacache_multiplier)
if callback != None:
callback(-1, None)
callback(-1, True)
for i, t in enumerate(tqdm(timesteps)):
latent_model_input = latents
slg_layers_local = None
@@ -281,7 +446,7 @@ class WanT2V:
# self.model.to(self.device)
if joint_pass:
noise_pred_cond, noise_pred_uncond = self.model(
latent_model_input, t=timestep,current_step=i, slg_layers=slg_layers_local, **arg_both)
latent_model_input, t=timestep, current_step=i, slg_layers=slg_layers_local, **arg_both)
if self._interrupt:
return None
else:
@@ -322,14 +487,18 @@ class WanT2V:
del temp_x0
if callback is not None:
callback(i, latents)
callback(i, False)
x0 = latents
if offload_model:
self.model.cpu()
torch.cuda.empty_cache()
if self.rank == 0:
videos = self.vae.decode(x0, VAE_tile_size)
if input_frames == None:
videos = self.vae.decode(x0, VAE_tile_size)
else:
videos = self.decode_latent(x0, input_ref_images, VAE_tile_size)
del noise, latents
@@ -341,3 +510,14 @@ class WanT2V:
dist.barrier()
return videos[0] if self.rank == 0 else None
def adapt_vace_model(self):
model = self.model
modules_dict= { k: m for k, m in model.named_modules()}
for model_layer, vace_layer in model.vace_layers_mapping.items():
module = modules_dict[f"vace_blocks.{vace_layer}"]
target = modules_dict[f"blocks.{model_layer}"]
setattr(target, "vace", module )
delattr(model, "vace_blocks")

80
wan/utils/utils.py
View File

@@ -3,14 +3,94 @@ import argparse
import binascii
import os
import os.path as osp
import torchvision.transforms.functional as TF
import torch.nn.functional as F
import imageio
import torch
import decord
import torchvision
from PIL import Image
import numpy as np
from rembg import remove, new_session
__all__ = ['cache_video', 'cache_image', 'str2bool']
from PIL import Image
def resample(video_fps, video_frames_count, max_target_frames_count, target_fps, start_target_frame ):
import math
video_frame_duration = 1 /video_fps
target_frame_duration = 1 / target_fps
target_time = start_target_frame * target_frame_duration
frame_no = math.ceil(target_time / video_frame_duration)
cur_time = frame_no * video_frame_duration
frame_ids =[]
while True:
if max_target_frames_count != 0 and len(frame_ids) >= max_target_frames_count :
break
add_frames_count = math.ceil( (target_time -cur_time) / video_frame_duration )
frame_no += add_frames_count
if frame_no >= video_frames_count:
break
frame_ids.append(frame_no)
cur_time += add_frames_count * video_frame_duration
target_time += target_frame_duration
frame_ids = frame_ids[:max_target_frames_count]
return frame_ids
def get_video_frame(file_name, frame_no):
decord.bridge.set_bridge('torch')
reader = decord.VideoReader(file_name)
frame = reader.get_batch([frame_no]).squeeze(0)
img = Image.fromarray(frame.numpy().astype(np.uint8))
return img
def resize_lanczos(img, h, w):
img = Image.fromarray(np.clip(255. * img.movedim(0, -1).cpu().numpy(), 0, 255).astype(np.uint8))
img = img.resize((w,h), resample=Image.Resampling.LANCZOS)
return torch.from_numpy(np.array(img).astype(np.float32) / 255.0).movedim(-1, 0)
def remove_background(img, session=None):
if session ==None:
session = new_session()
img = Image.fromarray(np.clip(255. * img.movedim(0, -1).cpu().numpy(), 0, 255).astype(np.uint8))
img = remove(img, session=session, alpha_matting = True, bgcolor=[255, 255, 255, 0]).convert('RGB')
return torch.from_numpy(np.array(img).astype(np.float32) / 255.0).movedim(-1, 0)
def resize_and_remove_background(img_list, canvas_width, canvas_height, rm_background ):
if rm_background:
session = new_session()
output_list =[]
for img in img_list:
width, height = img.size
white_canvas = np.full( (canvas_height, canvas_width, 3), 255, dtype= np.uint8 )
scale = min(canvas_height / height, canvas_width / width)
new_height = int(height * scale)
new_width = int(width * scale)
resized_image= img.resize((new_width,new_height), resample=Image.Resampling.LANCZOS)
if rm_background:
resized_image = remove(resized_image, session=session, alpha_matting = True, bgcolor=[255, 255, 255, 0]).convert('RGB')
top = (canvas_height - new_height) // 2
left = (canvas_width - new_width) // 2
white_canvas[top:top + new_height, left:left + new_width, :] = np.array(resized_image)
img = Image.fromarray(white_canvas)
output_list.append(img)
return output_list
def rand_name(length=8, suffix=''):
name = binascii.b2a_hex(os.urandom(length)).decode('utf-8')
if suffix:

303
wan/utils/vace_preprocessor.py Normal file
View File

@@ -0,0 +1,303 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import numpy as np
from PIL import Image
import torch
import torch.nn.functional as F
import torchvision.transforms.functional as TF
class VaceImageProcessor(object):
def __init__(self, downsample=None, seq_len=None):
self.downsample = downsample
self.seq_len = seq_len
def _pillow_convert(self, image, cvt_type='RGB'):
if image.mode != cvt_type:
if image.mode == 'P':
image = image.convert(f'{cvt_type}A')
if image.mode == f'{cvt_type}A':
bg = Image.new(cvt_type,
size=(image.width, image.height),
color=(255, 255, 255))
bg.paste(image, (0, 0), mask=image)
image = bg
else:
image = image.convert(cvt_type)
return image
def _load_image(self, img_path):
if img_path is None or img_path == '':
return None
img = Image.open(img_path)
img = self._pillow_convert(img)
return img
def _resize_crop(self, img, oh, ow, normalize=True):
"""
Resize, center crop, convert to tensor, and normalize.
"""
# resize and crop
iw, ih = img.size
if iw != ow or ih != oh:
# resize
scale = max(ow / iw, oh / ih)
img = img.resize(
(round(scale * iw), round(scale * ih)),
resample=Image.Resampling.LANCZOS
)
assert img.width >= ow and img.height >= oh
# center crop
x1 = (img.width - ow) // 2
y1 = (img.height - oh) // 2
img = img.crop((x1, y1, x1 + ow, y1 + oh))
# normalize
if normalize:
img = TF.to_tensor(img).sub_(0.5).div_(0.5).unsqueeze(1)
return img
def _image_preprocess(self, img, oh, ow, normalize=True, **kwargs):
return self._resize_crop(img, oh, ow, normalize)
def load_image(self, data_key, **kwargs):
return self.load_image_batch(data_key, **kwargs)
def load_image_pair(self, data_key, data_key2, **kwargs):
return self.load_image_batch(data_key, data_key2, **kwargs)
def load_image_batch(self, *data_key_batch, normalize=True, seq_len=None, **kwargs):
seq_len = self.seq_len if seq_len is None else seq_len
imgs = []
for data_key in data_key_batch:
img = self._load_image(data_key)
imgs.append(img)
w, h = imgs[0].size
dh, dw = self.downsample[1:]
# compute output size
scale = min(1., np.sqrt(seq_len / ((h / dh) * (w / dw))))
oh = int(h * scale) // dh * dh
ow = int(w * scale) // dw * dw
assert (oh // dh) * (ow // dw) <= seq_len
imgs = [self._image_preprocess(img, oh, ow, normalize) for img in imgs]
return *imgs, (oh, ow)
class VaceVideoProcessor(object):
def __init__(self, downsample, min_area, max_area, min_fps, max_fps, zero_start, seq_len, keep_last, **kwargs):
self.downsample = downsample
self.min_area = min_area
self.max_area = max_area
self.min_fps = min_fps
self.max_fps = max_fps
self.zero_start = zero_start
self.keep_last = keep_last
self.seq_len = seq_len
assert seq_len >= min_area / (self.downsample[1] * self.downsample[2])
@staticmethod
def resize_crop(video: torch.Tensor, oh: int, ow: int):
"""
Resize, center crop and normalize for decord loaded video (torch.Tensor type)
Parameters:
video - video to process (torch.Tensor): Tensor from `reader.get_batch(frame_ids)`, in shape of (T, H, W, C)
oh - target height (int)
ow - target width (int)
Returns:
The processed video (torch.Tensor): Normalized tensor range [-1, 1], in shape of (C, T, H, W)
Raises:
"""
# permute ([t, h, w, c] -> [t, c, h, w])
video = video.permute(0, 3, 1, 2)
# resize and crop
ih, iw = video.shape[2:]
if ih != oh or iw != ow:
# resize
scale = max(ow / iw, oh / ih)
video = F.interpolate(
video,
size=(round(scale * ih), round(scale * iw)),
mode='bicubic',
antialias=True
)
assert video.size(3) >= ow and video.size(2) >= oh
# center crop
x1 = (video.size(3) - ow) // 2
y1 = (video.size(2) - oh) // 2
video = video[:, :, y1:y1 + oh, x1:x1 + ow]
# permute ([t, c, h, w] -> [c, t, h, w]) and normalize
video = video.transpose(0, 1).float().div_(127.5).sub_(1.)
return video
def _video_preprocess(self, video, oh, ow):
return self.resize_crop(video, oh, ow)
def _get_frameid_bbox_default(self, fps, frame_timestamps, h, w, crop_box, rng):
target_fps = min(fps, self.max_fps)
duration = frame_timestamps[-1].mean()
x1, x2, y1, y2 = [0, w, 0, h] if crop_box is None else crop_box
h, w = y2 - y1, x2 - x1
ratio = h / w
df, dh, dw = self.downsample
# min/max area of the [latent video]
min_area_z = self.min_area / (dh * dw)
max_area_z = min(self.seq_len, self.max_area / (dh * dw), (h // dh) * (w // dw))
# sample a frame number of the [latent video]
rand_area_z = np.square(np.power(2, rng.uniform(
np.log2(np.sqrt(min_area_z)),
np.log2(np.sqrt(max_area_z))
)))
of = min(
(int(duration * target_fps) - 1) // df + 1,
int(self.seq_len / rand_area_z)
)
# deduce target shape of the [latent video]
target_area_z = min(max_area_z, int(self.seq_len / of))
oh = round(np.sqrt(target_area_z * ratio))
ow = int(target_area_z / oh)
of = (of - 1) * df + 1
oh *= dh
ow *= dw
# sample frame ids
target_duration = of / target_fps
begin = 0. if self.zero_start else rng.uniform(0, duration - target_duration)
timestamps = np.linspace(begin, begin + target_duration, of)
frame_ids = np.argmax(np.logical_and(
timestamps[:, None] >= frame_timestamps[None, :, 0],
timestamps[:, None] < frame_timestamps[None, :, 1]
), axis=1).tolist()
return frame_ids, (x1, x2, y1, y2), (oh, ow), target_fps
def _get_frameid_bbox_adjust_last(self, fps, video_frames_count, h, w, crop_box, rng, max_frames= 0, start_frame =0):
from wan.utils.utils import resample
target_fps = self.max_fps
# video_frames_count = len(frame_timestamps)
frame_ids= resample(fps, video_frames_count, max_frames, target_fps, start_frame )
x1, x2, y1, y2 = [0, w, 0, h] if crop_box is None else crop_box
h, w = y2 - y1, x2 - x1
ratio = h / w
df, dh, dw = self.downsample
seq_len = self.seq_len
# min/max area of the [latent video]
min_area_z = self.min_area / (dh * dw)
# max_area_z = min(seq_len, self.max_area / (dh * dw), (h // dh) * (w // dw))
max_area_z = min_area_z # workaround bug
# sample a frame number of the [latent video]
rand_area_z = np.square(np.power(2, rng.uniform(
np.log2(np.sqrt(min_area_z)),
np.log2(np.sqrt(max_area_z))
)))
seq_len = max_area_z * ((max_frames- start_frame - 1) // df +1)
# of = min(
# (len(frame_ids) - 1) // df + 1,
# int(seq_len / rand_area_z)
# )
of = (len(frame_ids) - 1) // df + 1
# deduce target shape of the [latent video]
# target_area_z = min(max_area_z, int(seq_len / of))
target_area_z = max_area_z
oh = round(np.sqrt(target_area_z * ratio))
ow = int(target_area_z / oh)
of = (of - 1) * df + 1
oh *= dh
ow *= dw
return frame_ids, (x1, x2, y1, y2), (oh, ow), target_fps
def _get_frameid_bbox(self, fps, video_frames_count, h, w, crop_box, rng, max_frames= 0, start_frame= 0):
if self.keep_last:
return self._get_frameid_bbox_adjust_last(fps, video_frames_count, h, w, crop_box, rng, max_frames= max_frames, start_frame= start_frame)
else:
return self._get_frameid_bbox_default(fps, video_frames_count, h, w, crop_box, rng, max_frames= max_frames)
def load_video(self, data_key, crop_box=None, seed=2024, **kwargs):
return self.load_video_batch(data_key, crop_box=crop_box, seed=seed, **kwargs)
def load_video_pair(self, data_key, data_key2, crop_box=None, seed=2024, **kwargs):
return self.load_video_batch(data_key, data_key2, crop_box=crop_box, seed=seed, **kwargs)
def load_video_batch(self, *data_key_batch, crop_box=None, seed=2024, max_frames= 0, trim_video =0, start_frame = 0, **kwargs):
rng = np.random.default_rng(seed + hash(data_key_batch[0]) % 10000)
# read video
import decord
decord.bridge.set_bridge('torch')
readers = []
src_video = None
for data_k in data_key_batch:
if torch.is_tensor(data_k):
src_video = data_k
else:
reader = decord.VideoReader(data_k)
readers.append(reader)
if src_video != None:
fps = 16
length = src_video.shape[0] + start_frame
if len(readers) > 0:
min_readers = min([len(r) for r in readers])
length = min(length, min_readers )
else:
fps = readers[0].get_avg_fps()
length = min([len(r) for r in readers])
# frame_timestamps = [readers[0].get_frame_timestamp(i) for i in range(length)]
# frame_timestamps = np.array(frame_timestamps, dtype=np.float32)
max_frames = min(max_frames, trim_video) if trim_video > 0 else max_frames
if src_video != None:
src_video = src_video[:max_frames]
h, w = src_video.shape[1:3]
else:
h, w = readers[0].next().shape[:2]
frame_ids, (x1, x2, y1, y2), (oh, ow), fps = self._get_frameid_bbox(fps, length, h, w, crop_box, rng, max_frames=max_frames, start_frame = start_frame )
# preprocess video
videos = [reader.get_batch(frame_ids)[:, y1:y2, x1:x2, :] for reader in readers]
if src_video != None:
videos = [src_video] + videos
videos = [self._video_preprocess(video, oh, ow) for video in videos]
return *videos, frame_ids, (oh, ow), fps
# return videos if len(videos) > 1 else videos[0]
def prepare_source(src_video, src_mask, src_ref_images, num_frames, image_size, device):
for i, (sub_src_video, sub_src_mask) in enumerate(zip(src_video, src_mask)):
if sub_src_video is None and sub_src_mask is None:
src_video[i] = torch.zeros((3, num_frames, image_size[0], image_size[1]), device=device)
src_mask[i] = torch.ones((1, num_frames, image_size[0], image_size[1]), device=device)
for i, ref_images in enumerate(src_ref_images):
if ref_images is not None:
for j, ref_img in enumerate(ref_images):
if ref_img is not None and ref_img.shape[-2:] != image_size:
canvas_height, canvas_width = image_size
ref_height, ref_width = ref_img.shape[-2:]
white_canvas = torch.ones((3, 1, canvas_height, canvas_width), device=device) # [-1, 1]
scale = min(canvas_height / ref_height, canvas_width / ref_width)
new_height = int(ref_height * scale)
new_width = int(ref_width * scale)
resized_image = F.interpolate(ref_img.squeeze(1).unsqueeze(0), size=(new_height, new_width), mode='bilinear', align_corners=False).squeeze(0).unsqueeze(1)
top = (canvas_height - new_height) // 2
left = (canvas_width - new_width) // 2
white_canvas[:, :, top:top + new_height, left:left + new_width] = resized_image
src_ref_images[i][j] = white_canvas
return src_video, src_mask, src_ref_images

4700
wgp.py Normal file
View File

File diff suppressed because it is too large Load Diff