Duplicate from nev/CoNR

Co-authored-by: Stepan Shabalin <[email protected]>

.gitattributes

@@ -0,0 +1,33 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+poses.zip filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,14 @@
+results/
+poses/
+test_data/
+test_data_pre/
+character_sheet/
+# weights/
+*.mp4
+*.webm
+poses.zip
+gradio_cached_examples/
+x264/
+filelist.txt
+complex_infer.sh
+__pycache__/

LICENSE

@@ -0,0 +1,7 @@
+Copyright 2022 Megvii Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

README.md

@@ -0,0 +1,127 @@
+---
+title: CoNR
+emoji: ⚡
+colorFrom: gray
+colorTo: red
+sdk: gradio
+sdk_version: 3.1.4
+app_file: app.py
+pinned: false
+license: mit
+duplicated_from: nev/CoNR
+---
+
+[English](https://github.com/megvii-research/CoNR/blob/main/README.md) | [中文](https://github.com/megvii-research/CoNR/blob/main/README_chinese.md)
+# Collaborative Neural Rendering using Anime Character Sheets 
+
+
+## [Homepage](https://conr.ml) | Colab [English](https://colab.research.google.com/github/megvii-research/CoNR/blob/main/notebooks/conr.ipynb)/[中文](https://colab.research.google.com/github/megvii-research/CoNR/blob/main/notebooks/conr_chinese.ipynb) | [arXiv](https://arxiv.org/abs/2207.05378)
+
+![image](images/MAIN.png)
+
+## Introduction
+
+This project is the official implement of [Collaborative Neural Rendering using Anime Character Sheets](https://arxiv.org/abs/2207.05378), which aims to genarate vivid dancing videos from hand-drawn anime character sheets(ACS). Watch more demos in our [HomePage](https://conr.ml).
+
+Contributors: [@transpchan](https://github.com/transpchan/), [@P2Oileen](https://github.com/P2Oileen), [@hzwer](https://github.com/hzwer)
+
+## Usage
+
+#### Prerequisites
+
+* NVIDIA GPU + CUDA + CUDNN
+* Python 3.6
+
+#### Installation
+
+* Clone this repository
+
+```bash
+git clone https://github.com/megvii-research/CoNR
+```
+
+* Dependencies
+
+To install all the dependencies, please run the following commands.
+
+```bash
+cd CoNR
+pip install -r requirements.txt
+```
+
+* Download Weights
+Download weights from Google Drive. Alternatively, you can download from [Baidu Netdisk](https://pan.baidu.com/s/1U11iIk-DiJodgCveSzB6ig?pwd=RDxc) (password:RDxc).
+
+```
+mkdir weights && cd weights
+gdown https://drive.google.com/uc?id=1M1LEpx70tJ72AIV2TQKr6NE_7mJ7tLYx
+gdown https://drive.google.com/uc?id=1YvZy3NHkJ6gC3pq_j8agcbEJymHCwJy0
+gdown https://drive.google.com/uc?id=1AOWZxBvTo9nUf2_9Y7Xe27ZFQuPrnx9i
+gdown https://drive.google.com/uc?id=19jM1-GcqgGoE1bjmQycQw_vqD9C5e-Jm
+```
+
+#### Prepare Inputs
+We provide two Ultra-Dense Pose sequences for two characters. You can generate more UDPs via 3D models and motions refers to [our paper](https://arxiv.org/abs/2207.05378).
+[Baidu Netdisk](https://pan.baidu.com/s/1hWvz4iQXnVTaTSb6vu1NBg?pwd=RDxc) (password:RDxc) 
+
+```
+# for short hair girl
+gdown https://drive.google.com/uc?id=11HMSaEkN__QiAZSnCuaM6GI143xo62KO
+unzip short_hair.zip
+mv short_hair/ poses/
+
+# for double ponytail girl
+gdown https://drive.google.com/uc?id=1WNnGVuU0ZLyEn04HzRKzITXqib1wwM4Q
+unzip double_ponytail.zip
+mv double_ponytail/ poses/
+```
+
+We provide sample inputs of anime character sheets. You can also draw more by yourself.
+Character sheets need to be cut out from the background and in png format.
+[Baidu Netdisk](https://pan.baidu.com/s/1shpP90GOMeHke7MuT0-Txw?pwd=RDxc) (password:RDxc) 
+
+```
+# for short hair girl
+gdown https://drive.google.com/uc?id=1r-3hUlENSWj81ve2IUPkRKNB81o9WrwT
+unzip short_hair_images.zip
+mv short_hair_images/ character_sheet/
+
+# for double ponytail girl
+gdown https://drive.google.com/uc?id=1XMrJf9Lk_dWgXyTJhbEK2LZIXL9G3MWc
+unzip double_ponytail_images.zip
+mv double_ponytail_images/ character_sheet/
+```
+
+#### RUN!
+* with web UI (powered by [Streamlit](https://streamlit.io/))
+
+```
+streamlit run streamlit.py --server.port=8501
+```
+then open your browser and visit `localhost:8501`, follow the instructions to genarate video.
+
+* via terminal
+
+```
+mkdir {dir_to_save_result}
+
+python -m torch.distributed.launch \
+--nproc_per_node=1 train.py --mode=test \
+--world_size=1 --dataloaders=2 \
+--test_input_poses_images={dir_to_poses} \
+--test_input_person_images={dir_to_character_sheet} \
+--test_output_dir={dir_to_save_result} \
+--test_checkpoint_dir={dir_to_weights}
+
+ffmpeg -r 30 -y -i {dir_to_save_result}/%d.png -r 30 -c:v libx264 output.mp4 -r 30
+```
+
+## Citation
+```bibtex
+@article{lin2022conr,
+  title={Collaborative Neural Rendering using Anime Character Sheets},
+  author={Lin, Zuzeng and Huang, Ailin and Huang, Zhewei and Hu, Chen and Zhou, Shuchang},
+  journal={arXiv preprint arXiv:2207.05378},
+  year={2022}
+}
+```

README_chinese.md

@@ -0,0 +1,124 @@
+[English](https://github.com/megvii-research/CoNR/blob/main/README.md) | [中文](https://github.com/megvii-research/CoNR/blob/main/README_chinese.md)
+# 用于二次元手绘设定稿动画化的神经渲染器
+
+
+## [HomePage](https://conr.ml) | Colab [English](https://colab.research.google.com/github/megvii-research/CoNR/blob/main/notebooks/conr.ipynb)/[中文](https://colab.research.google.com/github/megvii-research/CoNR/blob/main/notebooks/conr_chinese.ipynb) | [arXiv](https://arxiv.org/abs/2207.05378)
+
+![image](images/MAIN.png)
+
+## Introduction
+
+该项目为论文[Collaborative Neural Rendering using Anime Character Sheets](https://arxiv.org/abs/2207.05378)的官方复现，旨在从手绘人物设定稿生成生动的舞蹈动画。您可以在我们的[主页](https://conr.ml)中查看更多视频 demo。
+
+贡献者: [@transpchan](https://github.com/transpchan/), [@P2Oileen](https://github.com/P2Oileen), [@hzwer](https://github.com/hzwer)
+
+## 使用方法
+
+#### 需求
+
+* Nvidia GPU + CUDA + CUDNN
+* Python 3.6
+
+#### 安装
+
+* 克隆该项目
+
+```bash
+git clone https://github.com/megvii-research/CoNR
+```
+
+* 安装依赖
+
+请运行以下命令以安装CoNR所需的所有依赖。
+
+```bash
+cd CoNR
+pip install -r requirements.txt
+```
+
+* 下载权重
+运行以下代码，从 Google Drive 下载模型的权重。此外, 你也可以从 [百度云盘](https://pan.baidu.com/s/1U11iIk-DiJodgCveSzB6ig?pwd=RDxc) (password:RDxc)下载权重。
+
+```
+mkdir weights && cd weights
+gdown https://drive.google.com/uc?id=1M1LEpx70tJ72AIV2TQKr6NE_7mJ7tLYx
+gdown https://drive.google.com/uc?id=1YvZy3NHkJ6gC3pq_j8agcbEJymHCwJy0
+gdown https://drive.google.com/uc?id=1AOWZxBvTo9nUf2_9Y7Xe27ZFQuPrnx9i
+gdown https://drive.google.com/uc?id=19jM1-GcqgGoE1bjmQycQw_vqD9C5e-Jm
+```
+
+#### Prepare inputs
+我们为两个不同的人物，准备了两个超密集姿势（Ultra-Dense Pose）序列，从以下代码中二选一运行，即可从 Google Drive 下载。您可以通过任意的3D模型和动作数据，生成更多的超密集姿势序列，参考我们的[论文](https://arxiv.org/abs/2207.05378)。暂不提供官方转换接口。
+[百度云盘](https://pan.baidu.com/s/1hWvz4iQXnVTaTSb6vu1NBg?pwd=RDxc) (password:RDxc) 
+
+```
+# 短发女孩的超密集姿势
+gdown https://drive.google.com/uc?id=11HMSaEkN__QiAZSnCuaM6GI143xo62KO
+unzip short_hair.zip
+mv short_hair/ poses/
+
+# 双马尾女孩的超密集姿势
+gdown https://drive.google.com/uc?id=1WNnGVuU0ZLyEn04HzRKzITXqib1wwM4Q
+unzip double_ponytail.zip
+mv double_ponytail/ poses/
+```
+
+我们提供两个人物手绘设定表的样例，从以下代码中二选一运行，即可从 Google Drive下载。您也可以自行绘制。
+请注意：人物手绘设定表**必须从背景中分割开**，且必须为png格式。
+[百度云盘](https://pan.baidu.com/s/1shpP90GOMeHke7MuT0-Txw?pwd=RDxc) (password:RDxc) 
+
+```
+# 短发女孩的手绘设定表
+gdown https://drive.google.com/uc?id=1r-3hUlENSWj81ve2IUPkRKNB81o9WrwT
+unzip short_hair_images.zip
+mv short_hair_images/ character_sheet/
+
+# 双马尾女孩的手绘设定表
+gdown https://drive.google.com/uc?id=1XMrJf9Lk_dWgXyTJhbEK2LZIXL9G3MWc
+unzip double_ponytail_images.zip
+mv double_ponytail_images/ character_sheet/
+```
+
+#### 运行!
+我们提供两种方案：使用web图形界面，或使用命令行代码运行。
+
+* 使用web图形界面 (通过 [Streamlit](https://streamlit.io/) 实现)
+
+运行以下代码：
+
+```
+streamlit run streamlit.py --server_port.8501
+```
+
+然后打开浏览器并访问 `localhost:8501`, 根据页面内的指示生成视频。
+
+* 使用命令行代码
+
+请注意替换`{}`内容，并更换为您放置相应内容的文件夹位置。
+
+```
+mkdir {结果保存路径}
+
+python -m torch.distributed.launch \
+--nproc_per_node=1 train.py --mode=test \
+--world_size=1 --dataloaders=2 \
+--test_input_poses_images={姿势路径} \
+--test_input_person_images={人物设定表路径} \
+--test_output_dir={结果保存路径} \
+--test_checkpoint_dir={权重路径}
+
+ffmpeg -r 30 -y -i {结果保存路径}/%d.png -r 30 -c:v libx264 output.mp4 -r 30
+```
+
+视频结果将生成在 `CoNR/output.mp4`。
+
+## 引用CoNR
+```bibtex
+@article{lin2022conr,
+  title={Collaborative Neural Rendering using Anime Character Sheets},
+  author={Lin, Zuzeng and Huang, Ailin and Huang, Zhewei and Hu, Chen and Zhou, Shuchang},
+  journal={arXiv preprint arXiv:2207.05378},
+  year={2022}
+}
+```
+

app.py

@@ -0,0 +1,59 @@
+import numpy as np
+import gradio as gr
+import os
+import base64
+
+
+def get_base64(bin_file):
+    with open(bin_file, "rb") as f:
+        data = f.read()
+    return base64.b64encode(data).decode()
+
+
+def conr_fn(character_sheets, pose_zip):
+    os.system("rm character_sheet/*")
+    os.system("rm result/*")
+    os.system("rm poses/*")
+    os.makedirs("character_sheet", exist_ok=True)
+    for i, e in enumerate(character_sheets):
+        with open(f"character_sheet/{i}.png", "wb") as f:
+            e.seek(0)
+            f.write(e.read())
+            e.seek(0)
+    os.makedirs("poses", exist_ok=True)
+    pose_zip.seek(0)
+    open("poses.zip", "wb").write(pose_zip.read())
+    os.system(f"unzip -d poses poses.zip")
+    os.system("sh infer.sh")
+    return "output.mp4"
+
+
+with gr.Blocks() as ui:
+    gr.Markdown("CoNR demo")
+    gr.Markdown("<a target='_blank' href='https://colab.research.google.com/github/megvii-research/CoNR/blob/main/notebooks/conr.ipynb'> <img src='https://colab.research.google.com/assets/colab-badge.svg' alt='Open In Colab'/></a> [GitHub](https://github.com/megvii-research/CoNR/)")
+    gr.Markdown("Unofficial demo for [CoNR](https://transpchan.github.io/live3d/).")
+
+    with gr.Row():
+        # with gr.Column():
+            # gr.Markdown("## Parse video")   
+            # gr.Markdown("TBD")
+        with gr.Column():
+            gr.Markdown("## Animate character")
+            gr.Markdown("Character sheet")
+            character_sheets = gr.File(file_count="multiple")
+            gr.Markdown("Pose zip")  # Don't hack
+            pose_video = gr.File(file_count="single")
+
+            # os.system("sh download.sh")
+            run = gr.Button("Run")
+            video = gr.Video()
+            run.click(fn=conr_fn, inputs=[character_sheets, pose_video], outputs=video)
+            
+            gr.Markdown("## Examples")
+            sheets = "character_sheet_ponytail_example"
+            gr.Examples(fn=conr_fn, inputs=[character_sheets, pose_video], outputs=video,
+                        examples=[[[os.path.join(sheets, x) for x in os.listdir(sheets)], "poses_template.zip"]], cache_examples=True, examples_per_page=1)
+
+# ui.launch()
+demo = ui
+demo.launch()

conr.py

@@ -0,0 +1,292 @@
+
+import os
+import torch
+
+from model.backbone import ResEncUnet
+
+from model.shader import CINN
+from model.decoder_small import RGBADecoderNet
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def UDPClip(x):
+    return torch.clamp(x, min=0, max=1)  # NCHW
+
+
+class CoNR():
+    def __init__(self, args):
+        self.args = args
+
+        self.udpparsernet = ResEncUnet(
+            backbone_name='resnet50_danbo',
+            classes=4,
+            pretrained=(args.local_rank == 0),
+            parametric_upsampling=True,
+            decoder_filters=(512, 384, 256, 128, 32),
+            map_location=device
+        )
+        self.target_pose_encoder = ResEncUnet(
+            backbone_name='resnet18_danbo-4',
+            classes=1,
+            pretrained=(args.local_rank == 0),
+            parametric_upsampling=True,
+            decoder_filters=(512, 384, 256, 128, 32),
+            map_location=device
+        )
+        self.DIM_SHADER_REFERENCE = 4
+        self.shader = CINN(self.DIM_SHADER_REFERENCE)
+        self.rgbadecodernet = RGBADecoderNet(
+        )
+        self.device()
+        self.parser_ckpt = None
+
+    def dist(self):
+        args = self.args
+        if args.distributed:
+            self.udpparsernet = torch.nn.parallel.DistributedDataParallel(
+                self.udpparsernet,
+                device_ids=[
+                    args.local_rank],
+                output_device=args.local_rank,
+                broadcast_buffers=False,
+                find_unused_parameters=True
+            )
+            self.target_pose_encoder = torch.nn.parallel.DistributedDataParallel(
+                self.target_pose_encoder,
+                device_ids=[
+                    args.local_rank],
+                output_device=args.local_rank,
+                broadcast_buffers=False,
+                find_unused_parameters=True
+            )
+            self.shader = torch.nn.parallel.DistributedDataParallel(
+                self.shader,
+                device_ids=[
+                    args.local_rank],
+                output_device=args.local_rank,
+                broadcast_buffers=True
+            )
+
+            self.rgbadecodernet = torch.nn.parallel.DistributedDataParallel(
+                self.rgbadecodernet,
+                device_ids=[
+                    args.local_rank],
+                output_device=args.local_rank,
+                broadcast_buffers=True
+            )
+
+    def load_model(self, path):
+        self.udpparsernet.load_state_dict(
+            torch.load('{}/udpparsernet.pth'.format(path), map_location=device))
+        self.target_pose_encoder.load_state_dict(
+            torch.load('{}/target_pose_encoder.pth'.format(path), map_location=device))
+        self.shader.load_state_dict(
+            torch.load('{}/shader.pth'.format(path), map_location=device))
+        self.rgbadecodernet.load_state_dict(
+            torch.load('{}/rgbadecodernet.pth'.format(path), map_location=device))
+
+    def save_model(self, ite_num):
+        self._save_pth(self.udpparsernet,
+                       model_name="udpparsernet", ite_num=ite_num)
+        self._save_pth(self.target_pose_encoder,
+                       model_name="target_pose_encoder", ite_num=ite_num)
+        self._save_pth(self.shader,
+                       model_name="shader", ite_num=ite_num)
+        self._save_pth(self.rgbadecodernet,
+                       model_name="rgbadecodernet", ite_num=ite_num)
+
+    def _save_pth(self, net, model_name, ite_num):
+        args = self.args
+        to_save = None
+        if args.distributed:
+            if args.local_rank == 0:
+                to_save = net.module.state_dict()
+        else:
+            to_save = net.state_dict()
+        if to_save:
+            model_dir = os.path.join(
+                os.getcwd(), 'saved_models', args.model_name + os.sep + "checkpoints" + os.sep + "itr_%d" % (ite_num)+os.sep)
+
+            os.makedirs(model_dir, exist_ok=True)
+            torch.save(to_save, model_dir + model_name + ".pth")
+
+    def train(self):
+        self.udpparsernet.train()
+        self.target_pose_encoder.train()
+        self.shader.train()
+        self.rgbadecodernet.train()
+
+    def eval(self):
+        self.udpparsernet.eval()
+        self.target_pose_encoder.eval()
+        self.shader.eval()
+        self.rgbadecodernet.eval()
+
+    def device(self):
+        self.udpparsernet.to(device)
+        self.target_pose_encoder.to(device)
+        self.shader.to(device)
+        self.rgbadecodernet.to(device)
+
+    def data_norm_image(self, data):
+
+        with torch.cuda.amp.autocast(enabled=False):
+            for name in ["character_labels", "pose_label"]:
+                if name in data:
+                    data[name] = data[name].to(
+                        device, non_blocking=True).float()
+            for name in ["pose_images", "pose_mask", "character_images", "character_masks"]:
+                if name in data:
+                    data[name] = data[name].to(
+                        device, non_blocking=True).float() / 255.0
+            if "pose_images" in data:
+                data["num_pose_images"] = data["pose_images"].shape[1]
+                data["num_samples"] = data["pose_images"].shape[0]
+            if "character_images" in data:
+                data["num_character_images"] = data["character_images"].shape[1]
+                data["num_samples"] = data["character_images"].shape[0]
+            if "pose_images" in data and "character_images" in data:
+                assert (data["pose_images"].shape[0] ==
+                        data["character_images"].shape[0])
+        return data
+
+    def reset_charactersheet(self):
+        self.parser_ckpt = None
+
+    def model_step(self, data, training=False):
+        self.eval()
+        with torch.cuda.amp.autocast(enabled=False):
+            pred = {}
+            if self.parser_ckpt:
+                pred["parser"] = self.parser_ckpt
+            else:
+                pred = self.character_parser_forward(data, pred)
+                self.parser_ckpt = pred["parser"]
+            pred = self.pose_parser_sc_forward(data, pred)
+            pred = self.shader_pose_encoder_forward(data, pred)
+            pred = self.shader_forward(data, pred)
+        return pred
+
+    def shader_forward(self, data, pred={}):
+        assert ("num_character_images" in data), "ERROR: No Character Sheet input."
+
+        character_images_rgb_nmchw, num_character_images = data[
+            "character_images"], data["num_character_images"]
+        # build  x_reference_rgb_a_sudp in the draw call
+        shader_character_a_nmchw = data["character_masks"]
+        assert torch.any(torch.mean(shader_character_a_nmchw, (0, 2, 3, 4)) >= 0.95) == False, "ERROR: \
+                No transparent area found in the image, PLEASE separate the foreground of input character sheets.\
+                The website waifucutout.com is recommended to automatically cut out the foreground."
+        
+        if shader_character_a_nmchw is None:
+            shader_character_a_nmchw = pred["parser"]["pred"][:, :, 3:4, :, :]
+        x_reference_rgb_a = torch.cat([shader_character_a_nmchw[:, :, :, :, :] * character_images_rgb_nmchw[:, :, :, :, :],
+                                       shader_character_a_nmchw[:,
+                                                                :, :, :, :],
+
+                                       ], 2)
+        assert (x_reference_rgb_a.shape[2] == self.DIM_SHADER_REFERENCE)
+        # build  x_reference_features in the draw call
+        x_reference_features = pred["parser"]["features"]
+        # run cinn shader
+        retdic = self.shader(
+            pred["shader"]["target_pose_features"], x_reference_rgb_a, x_reference_features)
+        pred["shader"].update(retdic)
+
+        # decode rgba
+        if True:
+            dec_out = self.rgbadecodernet(
+                retdic["y_last_remote_features"])
+            y_weighted_x_reference_RGB = dec_out[:, 0:3, :, :]
+            y_weighted_mask_A = dec_out[:, 3:4, :, :]
+        y_weighted_warp_decoded_rgba = torch.cat(
+            (y_weighted_x_reference_RGB*y_weighted_mask_A, y_weighted_mask_A), dim=1
+        )
+        assert(y_weighted_warp_decoded_rgba.shape[1] == 4)
+        assert(
+            y_weighted_warp_decoded_rgba.shape[-1] == character_images_rgb_nmchw.shape[-1])
+        # apply decoded mask to decoded rgb, finishing the draw call
+        pred["shader"]["y_weighted_warp_decoded_rgba"] = y_weighted_warp_decoded_rgba
+        return pred
+
+    def character_parser_forward(self, data, pred={}):
+        if not("num_character_images" in data and "character_images" in data):
+            return pred
+        pred["parser"] = {"pred": None}  # create output
+
+        inputs_rgb_nmchw, num_samples, num_character_images = data[
+            "character_images"],  data["num_samples"], data["num_character_images"]
+        inputs_rgb_fchw = inputs_rgb_nmchw.view(
+            (num_samples * num_character_images, inputs_rgb_nmchw.shape[2], inputs_rgb_nmchw.shape[3], inputs_rgb_nmchw.shape[4]))
+
+        encoder_out, features = self.udpparsernet(
+            (inputs_rgb_fchw-0.6)/0.2970)
+
+        pred["parser"]["features"] = [features_out.view(
+            (num_samples, num_character_images, features_out.shape[1], features_out.shape[2], features_out.shape[3])) for features_out in features]
+
+        if (encoder_out is not None):
+
+            pred["parser"]["pred"] = UDPClip(encoder_out.view(
+                (num_samples, num_character_images, encoder_out.shape[1], encoder_out.shape[2], encoder_out.shape[3])))
+
+        return pred
+
+    def pose_parser_sc_forward(self, data, pred={}):
+        if not("num_pose_images" in data and "pose_images" in data):
+            return pred
+        inputs_aug_rgb_nmchw, num_samples, num_pose_images = data[
+            "pose_images"],  data["num_samples"], data["num_pose_images"]
+        inputs_aug_rgb_fchw = inputs_aug_rgb_nmchw.view(
+            (num_samples * num_pose_images, inputs_aug_rgb_nmchw.shape[2], inputs_aug_rgb_nmchw.shape[3], inputs_aug_rgb_nmchw.shape[4]))
+
+        encoder_out, _ = self.udpparsernet(
+            (inputs_aug_rgb_fchw-0.6)/0.2970)
+
+        encoder_out = encoder_out.view(
+            (num_samples, num_pose_images, encoder_out.shape[1], encoder_out.shape[2], encoder_out.shape[3]))
+
+        # apply sigmoid after eval loss
+        pred["pose_parser"] = {"pred":UDPClip(encoder_out)[:,0,:,:,:]}
+        
+
+        return pred
+
+    def shader_pose_encoder_forward(self, data, pred={}):
+        pred["shader"] = {}  # create output
+        if "pose_images" in data:
+            pose_images_rgb_nmchw = data["pose_images"]
+            target_gt_rgb = pose_images_rgb_nmchw[:, 0, :, :, :]
+            pred["shader"]["target_gt_rgb"] = target_gt_rgb
+
+        shader_target_a = None
+        if "pose_mask" in data:
+            pred["shader"]["target_gt_a"] = data["pose_mask"]
+            shader_target_a = data["pose_mask"]
+
+        shader_target_sudp = None
+        if "pose_label" in data:
+            shader_target_sudp = data["pose_label"][:, :3, :, :]
+
+        if self.args.test_pose_use_parser_udp:
+            shader_target_sudp = None
+        if shader_target_sudp is None:
+            shader_target_sudp = pred["pose_parser"]["pred"][:, 0:3, :, :]
+
+        if shader_target_a is None:
+            shader_target_a = pred["pose_parser"]["pred"][:, 3:4, :, :]
+
+        # build x_target_sudp_a in the draw call
+        x_target_sudp_a = torch.cat((
+            shader_target_sudp*shader_target_a,
+            shader_target_a
+        ), 1)
+        pred["shader"].update({
+            "x_target_sudp_a": x_target_sudp_a
+        })
+        _, features = self.target_pose_encoder(
+            (x_target_sudp_a-0.6)/0.2970, ret_parser_out=False)
+
+        pred["shader"]["target_pose_features"] = features
+        return pred

data_loader.py

@@ -0,0 +1,275 @@
+import cv2
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+import os
+cv2.setNumThreads(1)
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
+
+
+class RandomResizedCropWithAutoCenteringAndZeroPadding (object):
+    def __init__(self, output_size, scale=(0.08, 1.0), ratio=(3. / 4., 4. / 3.), center_jitter=(0.1, 0.1), size_from_alpha_mask=True):
+        assert isinstance(output_size, (int, tuple))
+        if isinstance(output_size, int):
+            self.output_size = (output_size, output_size)
+        else:
+            assert len(output_size) == 2
+            self.output_size = output_size
+        assert isinstance(scale,  tuple)
+        assert isinstance(ratio,  tuple)
+        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
+            raise ValueError("Scale and ratio should be of kind (min, max)")
+        self.size_from_alpha_mask = size_from_alpha_mask
+        self.scale = scale
+        self.ratio = ratio
+        assert isinstance(center_jitter,  tuple)
+        self.center_jitter = center_jitter
+
+    def __call__(self, sample):
+        imidx, image = sample['imidx'], sample["image_np"]
+        if "labels" in sample:
+            label = sample["labels"]
+        else:
+            label = None
+
+        im_h, im_w = image.shape[:2]
+        if self.size_from_alpha_mask and image.shape[2] == 4:
+            # compute bbox from alpha mask
+            bbox_left, bbox_top, bbox_w, bbox_h = cv2.boundingRect(
+                (image[:, :, 3] > 0).astype(np.uint8))
+        else:
+            bbox_left, bbox_top = 0, 0
+            bbox_h, bbox_w = image.shape[:2]
+        if bbox_h <= 1 and bbox_w <= 1:
+            sample["bad"] = 0
+        else:
+            # detect too small image here
+            alpha_varea = np.sum((image[:, :, 3] > 0).astype(np.uint8))
+            image_area = image.shape[0]*image.shape[1]
+            if alpha_varea/image_area < 0.001:
+                sample["bad"] = alpha_varea
+                # detect bad image
+        if "bad" in sample:
+            # baddata_dir = os.path.join(os.getcwd(), 'test_data', "baddata" + os.sep)
+            # save_output(str(imidx)+".png",image,label,baddata_dir)
+            bbox_h, bbox_w = image.shape[:2]
+            sample["image_np"] = np.zeros(
+                [self.output_size[0], self.output_size[1], image.shape[2]], dtype=image.dtype)
+            if label is not None:
+                sample["labels"] = np.zeros(
+                    [self.output_size[0], self.output_size[1], 4], dtype=label.dtype)
+
+            return sample
+
+        # compute default area by making sure output_size contains bbox_w * bbox_h
+
+        jitter_h = np.random.uniform(-bbox_h *
+                                     self.center_jitter[0], bbox_h*self.center_jitter[0])
+        jitter_w = np.random.uniform(-bbox_w *
+                                     self.center_jitter[1], bbox_w*self.center_jitter[1])
+
+        # h/w
+        target_aspect_ratio = np.exp(
+            np.log(self.output_size[0]/self.output_size[1]) +
+            np.random.uniform(np.log(self.ratio[0]), np.log(self.ratio[1]))
+        )
+
+        source_aspect_ratio = bbox_h/bbox_w
+
+        if target_aspect_ratio < source_aspect_ratio:
+            # same w, target has larger h, use h to align
+            target_height = bbox_h * \
+                np.random.uniform(self.scale[0], self.scale[1])
+            virtual_h = int(
+                round(target_height))
+            virtual_w = int(
+                round(target_height / target_aspect_ratio))  # h/w
+        else:
+            # same w, source has larger h, use w to align
+            target_width = bbox_w * \
+                np.random.uniform(self.scale[0], self.scale[1])
+            virtual_h = int(
+                round(target_width * target_aspect_ratio))  # h/w
+            virtual_w = int(
+                round(target_width))
+
+        # print("required aspect ratio:",      target_aspect_ratio)
+
+        virtual_top = int(round(bbox_top + jitter_h - (virtual_h-bbox_h)/2))
+        virutal_left = int(round(bbox_left + jitter_w - (virtual_w-bbox_w)/2))
+
+        if virtual_top < 0:
+            top_padding = abs(virtual_top)
+            crop_top = 0
+        else:
+            top_padding = 0
+            crop_top = virtual_top
+        if virutal_left < 0:
+            left_padding = abs(virutal_left)
+            crop_left = 0
+        else:
+            left_padding = 0
+            crop_left = virutal_left
+        if virtual_top+virtual_h > im_h:
+            bottom_padding = abs(im_h-(virtual_top+virtual_h))
+            crop_bottom = im_h
+        else:
+            bottom_padding = 0
+            crop_bottom = virtual_top+virtual_h
+        if virutal_left+virtual_w > im_w:
+            right_padding = abs(im_w-(virutal_left+virtual_w))
+            crop_right = im_w
+        else:
+            right_padding = 0
+            crop_right = virutal_left+virtual_w
+        # crop
+
+        image = image[crop_top:crop_bottom, crop_left:  crop_right]
+        if label is not None:
+            label = label[crop_top:crop_bottom, crop_left:  crop_right]
+
+        # pad
+        if top_padding + bottom_padding + left_padding + right_padding > 0:
+            padding = ((top_padding, bottom_padding),
+                       (left_padding, right_padding), (0, 0))
+            # print("padding", padding)
+            image = np.pad(image, padding, mode='constant')
+            if label is not None:
+                label = np.pad(label, padding, mode='constant')
+
+        if image.shape[0]/image.shape[1] - virtual_h/virtual_w > 0.001:
+            print("virtual aspect ratio:",  virtual_h/virtual_w)
+            print("image aspect ratio:", image.shape[0]/image.shape[1])
+        assert (image.shape[0]/image.shape[1] - virtual_h/virtual_w < 0.001)
+        sample["crop"] = np.array(
+            [im_h, im_w,  crop_top, crop_bottom, crop_left, crop_right, top_padding, bottom_padding, left_padding, right_padding, image.shape[0], image.shape[1]])
+
+        # resize
+        if self.output_size[1] != image.shape[1] or self.output_size[0] != image.shape[0]:
+            if self.output_size[1] > image.shape[1] and self.output_size[0] > image.shape[0]:
+                # enlarging
+                image = cv2.resize(
+                    image, (self.output_size[1], self.output_size[0]), interpolation=cv2.INTER_LINEAR)
+            else:
+                # shrinking
+                image = cv2.resize(
+                    image, (self.output_size[1], self.output_size[0]), interpolation=cv2.INTER_AREA)
+
+            if label is not None:
+                label = cv2.resize(label, (self.output_size[1], self.output_size[0]),
+                                   interpolation=cv2.INTER_NEAREST_EXACT)
+
+        assert image.shape[0] == self.output_size[0] and image.shape[1] == self.output_size[1]
+        sample['imidx'], sample["image_np"] = imidx, image
+        if label is not None:
+            assert label.shape[0] == self.output_size[0] and label.shape[1] == self.output_size[1]
+            sample["labels"] = label
+
+        return sample
+
+
+class FileDataset(Dataset):
+    def __init__(self, image_names_list, fg_img_lbl_transform=None, shader_pose_use_gt_udp_test=True, shader_target_use_gt_rgb_debug=False):
+        self.image_name_list = image_names_list
+        self.fg_img_lbl_transform = fg_img_lbl_transform
+        self.shader_pose_use_gt_udp_test = shader_pose_use_gt_udp_test
+        self.shader_target_use_gt_rgb_debug = shader_target_use_gt_rgb_debug
+
+    def __len__(self):
+        return len(self.image_name_list)
+
+    def get_gt_from_disk(self, idx, imname, read_label):
+        if read_label:
+            # read label
+            with open(imname, mode="rb") as bio:
+                if imname.find(".npz") > 0:
+                    label_np = np.load(bio, allow_pickle=True)[
+                        'i'].astype(np.float32, copy=False)
+                else:
+                    label_np = cv2.cvtColor(cv2.imdecode(np.frombuffer(bio.read(
+                    ), np.uint8),  cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH | cv2.IMREAD_UNCHANGED), cv2.COLOR_BGRA2RGBA)
+            assert (4 == label_np.shape[2])
+            # fake image out of valid label
+            image_np = (label_np*255).clip(0, 255).astype(np.uint8, copy=False)
+            # assemble sample
+            sample = {'imidx': np.array(
+                [idx]),   "image_np": image_np, "labels": label_np}
+
+        else:
+            # read image as unit8
+            with open(imname, mode="rb") as bio:
+                image_np = cv2.cvtColor(cv2.imdecode(np.frombuffer(
+                    bio.read(), np.uint8), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGRA2RGBA)
+                # image_np = Image.open(bio)
+                # image_np = np.array(image_np)
+            assert (3 == len(image_np.shape))
+            if (image_np.shape[2] == 4):
+                mask_np = image_np[:, :, 3:4]
+                image_np = (image_np[:, :, :3] *
+                            (image_np[:, :, 3][:, :, np.newaxis]/255.0)).clip(0, 255).astype(np.uint8, copy=False)
+            elif (image_np.shape[2] == 3):
+                # generate a fake mask
+                # Fool-proofing
+                mask_np = np.ones(
+                    (image_np.shape[0], image_np.shape[1], 1), dtype=np.uint8)*255
+                print("WARN: transparent background is preferred for image ", imname)
+            else:
+                raise ValueError("weird shape of image ", imname, image_np)
+            image_np = np.concatenate((image_np, mask_np), axis=2)
+            sample = {'imidx': np.array(
+                [idx]),   "image_np": image_np}
+
+        # apply fg_img_lbl_transform
+        if self.fg_img_lbl_transform:
+            sample = self.fg_img_lbl_transform(sample)
+
+        if "labels" in sample:
+            # return UDP as 4chn XYZV float tensor
+            if "float" not in str(sample["labels"].dtype):
+                sample["labels"] = sample["labels"].astype(np.float32) / np.iinfo(sample["labels"].dtype).max
+            sample["labels"] = torch.from_numpy(
+                sample["labels"].transpose((2, 0, 1)))
+            assert (sample["labels"].dtype == torch.float32)
+
+        if "image_np" in sample:
+            # return image as 3chn RGB uint8 tensor and 1chn A uint8 tensor
+            sample["mask"] = torch.from_numpy(
+                sample["image_np"][:, :, 3:4].transpose((2, 0, 1)))
+            assert (sample["mask"].dtype == torch.uint8)
+            sample["image"] = torch.from_numpy(
+                sample["image_np"][:, :, :3].transpose((2, 0, 1)))
+
+            assert (sample["image"].dtype == torch.uint8)
+            del sample["image_np"]
+        return sample
+
+    def __getitem__(self, idx):
+        sample = {
+            'imidx': np.array([idx])}
+        target = self.get_gt_from_disk(
+            idx, imname=self.image_name_list[idx][0], read_label=self.shader_pose_use_gt_udp_test)
+        if self.shader_target_use_gt_rgb_debug:
+            sample["pose_images"] = torch.stack([target["image"]])
+            sample["pose_mask"] = target["mask"]
+        elif self.shader_pose_use_gt_udp_test:
+            sample["pose_label"] = target["labels"]
+            sample["pose_mask"] = target["mask"]
+        else:
+            sample["pose_images"] = torch.stack([target["image"]])
+        if "crop" in target:
+            sample["pose_crop"] = target["crop"]
+        character_images = []
+        character_masks = []
+        for i in range(1, len(self.image_name_list[idx])):
+            source = self.get_gt_from_disk(
+                idx, self.image_name_list[idx][i], read_label=False)
+            character_images.append(source["image"])
+            character_masks.append(source["mask"])
+        character_images = torch.stack(character_images)
+        character_masks = torch.stack(character_masks)
+        sample.update({
+            "character_images": character_images,
+            "character_masks": character_masks
+        })
+        # do not make fake labels in inference
+        return sample

infer.sh

@@ -0,0 +1,13 @@
+rm -r "./results"
+mkdir "./results"
+
+python3 train.py --mode=test \
+--world_size=1 --dataloaders=2 \
+--test_input_poses_images=./poses/ \
+--test_input_person_images=./character_sheet/ \
+--test_output_dir=./results/ \
+--test_checkpoint_dir=./weights/ 
+
+echo Generating Video...
+ffmpeg -r 30 -y -i ./results/%d.png -r 30 -c:v libx264 -pix_fmt yuv420p output.mp4 -crf 18 -r 30
+echo DONE.

model/__init__.py

@@ -0,0 +1 @@
+

model/backbone.py

@@ -0,0 +1,289 @@
+"""
+This code was mostly taken from backbone-unet by mkisantal:
+https://github.com/mkisantal/backboned-unet/blob/master/backboned_unet/unet.py
+"""
+import torch
+import torch.nn as nn
+from torchvision import models
+from torch.nn import functional as F
+
+import torch.nn as nn
+import torch
+from torchvision import models
+
+
+class AdaptiveConcatPool2d(nn.Module):
+    """
+    Layer that concats `AdaptiveAvgPool2d` and `AdaptiveMaxPool2d`.
+    Source: Fastai. This code was taken from the fastai library at url
+    https://github.com/fastai/fastai/blob/master/fastai/layers.py#L176
+    """
+
+    def __init__(self, sz=None):
+        "Output will be 2*sz or 2 if sz is None"
+        super().__init__()
+        self.output_size = sz or 1
+        self.ap = nn.AdaptiveAvgPool2d(self.output_size)
+        self.mp = nn.AdaptiveMaxPool2d(self.output_size)
+
+    def forward(self, x): return torch.cat([self.mp(x), self.ap(x)], 1)
+
+
+class MyNorm(nn.Module):
+    def __init__(self, num_channels):
+        super(MyNorm, self).__init__()
+        self.norm = nn.InstanceNorm2d(
+            num_channels, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
+
+    def forward(self, x):
+        x = self.norm(x)
+        return x
+
+
+def resnet_fastai(model, pretrained,  url, replace_first_layer=None, replace_maxpool_layer=None, progress=True, map_location=None, **kwargs):
+    cut = -2
+    s = model(pretrained=False, **kwargs)
+    if replace_maxpool_layer is not None:
+        s.maxpool = replace_maxpool_layer
+    if replace_first_layer is not None:
+        body = nn.Sequential(replace_first_layer, *list(s.children())[1:cut])
+    else:
+        body = nn.Sequential(*list(s.children())[:cut])
+
+    if pretrained:
+        state = torch.hub.load_state_dict_from_url(url,
+                                                   progress=progress, map_location=map_location)
+        if replace_first_layer is not None:
+            for each in list(state.keys()).copy():
+                if each.find("0.0.") == 0:
+                    del state[each]
+        body_tail = nn.Sequential(body)
+        ret = body_tail.load_state_dict(state, strict=False)
+    return body
+
+
+def get_backbone(name, pretrained=True, map_location=None):
+    """ Loading backbone, defining names for skip-connections and encoder output. """
+
+    first_layer_for_4chn = nn.Conv2d(
+        4, 64, kernel_size=7, stride=2, padding=3, bias=False)
+    max_pool_layer_replace = nn.Conv2d(
+        64, 64, kernel_size=3, stride=2, padding=1, bias=False)
+    # loading backbone model
+    if name == 'resnet18':
+        backbone = models.resnet18(pretrained=pretrained)
+    if name == 'resnet18-4':
+        backbone = models.resnet18(pretrained=pretrained)
+        backbone.conv1 = first_layer_for_4chn
+    elif name == 'resnet34':
+        backbone = models.resnet34(pretrained=pretrained)
+    elif name == 'resnet50':
+        backbone = models.resnet50(pretrained=False, norm_layer=MyNorm)
+        backbone.maxpool = max_pool_layer_replace
+    elif name == 'resnet101':
+        backbone = models.resnet101(pretrained=pretrained)
+    elif name == 'resnet152':
+        backbone = models.resnet152(pretrained=pretrained)
+    elif name == 'vgg16':
+        backbone = models.vgg16_bn(pretrained=pretrained).features
+    elif name == 'vgg19':
+        backbone = models.vgg19_bn(pretrained=pretrained).features
+    elif name == 'resnet18_danbo-4':
+        backbone = resnet_fastai(models.resnet18, url="https://github.com/RF5/danbooru-pretrained/releases/download/v0.1/resnet18-3f77756f.pth",
+                                 pretrained=pretrained, map_location=map_location, norm_layer=MyNorm, replace_first_layer=first_layer_for_4chn)
+    elif name == 'resnet50_danbo':
+        backbone = resnet_fastai(models.resnet50, url="https://github.com/RF5/danbooru-pretrained/releases/download/v0.1/resnet50-13306192.pth",
+                                 pretrained=pretrained, map_location=map_location, norm_layer=MyNorm, replace_maxpool_layer=max_pool_layer_replace)
+    elif name == 'densenet121':
+        backbone = models.densenet121(pretrained=True).features
+    elif name == 'densenet161':
+        backbone = models.densenet161(pretrained=True).features
+    elif name == 'densenet169':
+        backbone = models.densenet169(pretrained=True).features
+    elif name == 'densenet201':
+        backbone = models.densenet201(pretrained=True).features
+    else:
+        raise NotImplemented(
+            '{} backbone model is not implemented so far.'.format(name))
+    #print(backbone)
+    # specifying skip feature and output names
+    if name.startswith('resnet'):
+        feature_names = [None, 'relu', 'layer1', 'layer2', 'layer3']
+        backbone_output = 'layer4'
+    elif name == 'vgg16':
+        # TODO: consider using a 'bridge' for VGG models, there is just a MaxPool between last skip and backbone output
+        feature_names = ['5', '12', '22', '32', '42']
+        backbone_output = '43'
+    elif name == 'vgg19':
+        feature_names = ['5', '12', '25', '38', '51']
+        backbone_output = '52'
+    elif name.startswith('densenet'):
+        feature_names = [None, 'relu0', 'denseblock1',
+                         'denseblock2', 'denseblock3']
+        backbone_output = 'denseblock4'
+    elif name == 'unet_encoder':
+        feature_names = ['module1', 'module2', 'module3', 'module4']
+        backbone_output = 'module5'
+    else:
+        raise NotImplemented(
+            '{} backbone model is not implemented so far.'.format(name))
+    if name.find('_danbo') > 0:
+        feature_names = [None, '2', '4', '5', '6']
+        backbone_output = '7'
+    return backbone, feature_names, backbone_output
+
+
+class UpsampleBlock(nn.Module):
+
+    # TODO: separate parametric and non-parametric classes?
+    # TODO: skip connection concatenated OR added
+
+    def __init__(self, ch_in, ch_out=None, skip_in=0, use_bn=True, parametric=False):
+        super(UpsampleBlock, self).__init__()
+
+        self.parametric = parametric
+        ch_out = ch_in/2 if ch_out is None else ch_out
+
+        # first convolution: either transposed conv, or conv following the skip connection
+        if parametric:
+            # versions: kernel=4 padding=1, kernel=2 padding=0
+            self.up = nn.ConvTranspose2d(in_channels=ch_in, out_channels=ch_out, kernel_size=(4, 4),
+                                         stride=2, padding=1, output_padding=0, bias=(not use_bn))
+            self.bn1 = MyNorm(ch_out) if use_bn else None
+        else:
+            self.up = None
+            ch_in = ch_in + skip_in
+            self.conv1 = nn.Conv2d(in_channels=ch_in, out_channels=ch_out, kernel_size=(3, 3),
+                                   stride=1, padding=1, bias=(not use_bn))
+            self.bn1 = MyNorm(ch_out) if use_bn else None
+
+        self.relu = nn.ReLU(inplace=True)
+
+        # second convolution
+        conv2_in = ch_out if not parametric else ch_out + skip_in
+        self.conv2 = nn.Conv2d(in_channels=conv2_in, out_channels=ch_out, kernel_size=(3, 3),
+                               stride=1, padding=1, bias=(not use_bn))
+        self.bn2 = MyNorm(ch_out) if use_bn else None
+
+    def forward(self, x, skip_connection=None):
+
+        x = self.up(x) if self.parametric else F.interpolate(x, size=None, scale_factor=2, mode='bilinear',
+                                                             align_corners=None)
+        if self.parametric:
+            x = self.bn1(x) if self.bn1 is not None else x
+            x = self.relu(x)
+
+        if skip_connection is not None:
+            x = torch.cat([x, skip_connection], dim=1)
+
+        if not self.parametric:
+            x = self.conv1(x)
+            x = self.bn1(x) if self.bn1 is not None else x
+            x = self.relu(x)
+        x = self.conv2(x)
+        x = self.bn2(x) if self.bn2 is not None else x
+        x = self.relu(x)
+
+        return x
+
+
+class ResEncUnet(nn.Module):
+
+    """ U-Net (https://arxiv.org/pdf/1505.04597.pdf) implementation with pre-trained torchvision backbones."""
+
+    def __init__(self,
+                 backbone_name,
+                 pretrained=True,
+                 encoder_freeze=False,
+                 classes=21,
+                 decoder_filters=(512, 256, 128, 64, 32),
+                 parametric_upsampling=True,
+                 shortcut_features='default',
+                 decoder_use_instancenorm=True,
+                 map_location=None
+                 ):
+        super(ResEncUnet, self).__init__()
+
+        self.backbone_name = backbone_name
+
+        self.backbone, self.shortcut_features, self.bb_out_name = get_backbone(
+            backbone_name, pretrained=pretrained, map_location=map_location)
+        shortcut_chs, bb_out_chs = self.infer_skip_channels()
+        if shortcut_features != 'default':
+            self.shortcut_features = shortcut_features
+
+        # build decoder part
+        self.upsample_blocks = nn.ModuleList()
+        # avoiding having more blocks than skip connections
+        decoder_filters = decoder_filters[:len(self.shortcut_features)]
+        decoder_filters_in = [bb_out_chs] + list(decoder_filters[:-1])
+        num_blocks = len(self.shortcut_features)
+        for i, [filters_in, filters_out] in enumerate(zip(decoder_filters_in, decoder_filters)):
+            self.upsample_blocks.append(UpsampleBlock(filters_in, filters_out,
+                                                      skip_in=shortcut_chs[num_blocks-i-1],
+                                                      parametric=parametric_upsampling,
+                                                      use_bn=decoder_use_instancenorm))
+        self.final_conv = nn.Conv2d(
+            decoder_filters[-1], classes, kernel_size=(1, 1))
+
+        if encoder_freeze:
+            self.freeze_encoder()
+
+    def freeze_encoder(self):
+        """ Freezing encoder parameters, the newly initialized decoder parameters are remaining trainable. """
+
+        for param in self.backbone.parameters():
+            param.requires_grad = False
+
+    def forward(self, *input, ret_parser_out=True):
+        """ Forward propagation in U-Net. """
+
+        x, features = self.forward_backbone(*input)
+        output_feature = [x]
+        for skip_name, upsample_block in zip(self.shortcut_features[::-1], self.upsample_blocks):
+            skip_features = features[skip_name]
+            if skip_features is not None:
+                output_feature.append(skip_features)
+            if ret_parser_out:
+                x = upsample_block(x, skip_features)
+        if ret_parser_out:
+            x = self.final_conv(x)
+            # apply sigmoid later
+        else:
+            x = None
+
+        return x, output_feature
+
+    def forward_backbone(self, x):
+        """ Forward propagation in backbone encoder network.  """
+
+        features = {None: None} if None in self.shortcut_features else dict()
+        for name, child in self.backbone.named_children():
+            x = child(x)
+            if name in self.shortcut_features:
+                features[name] = x
+            if name == self.bb_out_name:
+                break
+
+        return x, features
+
+    def infer_skip_channels(self):
+        """ Getting the number of channels at skip connections and at the output of the encoder. """
+        if self.backbone_name.find("-4") > 0:
+            x = torch.zeros(1, 4, 224, 224)
+        else:
+            x = torch.zeros(1, 3, 224, 224)
+        has_fullres_features = self.backbone_name.startswith(
+            'vgg') or self.backbone_name == 'unet_encoder'
+        # only VGG has features at full resolution
+        channels = [] if has_fullres_features else [0]
+
+        # forward run in backbone to count channels (dirty solution but works for *any* Module)
+        for name, child in self.backbone.named_children():
+            x = child(x)
+            if name in self.shortcut_features:
+                channels.append(x.shape[1])
+            if name == self.bb_out_name:
+                out_channels = x.shape[1]
+                break
+        return channels, out_channels

model/decoder_small.py

@@ -0,0 +1,43 @@
+
+
+from torch import nn
+import torch.nn.functional as F
+import torch
+
+
+class ResBlock2d(nn.Module):
+    def __init__(self, in_features, kernel_size, padding):
+        super(ResBlock2d, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size,
+                               padding=padding)
+        self.conv2 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size,
+                               padding=padding)
+    
+        self.norm1 = nn.Conv2d(
+            in_channels=in_features, out_channels=in_features, kernel_size=1)
+        self.norm2 = nn.Conv2d(
+            in_channels=in_features, out_channels=in_features, kernel_size=1)
+
+    def forward(self, x):
+        out = self.norm1(x)
+        out = F.relu(out, inplace=True)
+        out = self.conv1(out)
+        out = self.norm2(out)
+        out = F.relu(out, inplace=True)
+        out = self.conv2(out)
+        out += x
+        return out
+
+
+class RGBADecoderNet(nn.Module):
+    def __init__(self,  c=64, out_planes=4,  num_bottleneck_blocks=1):
+        super(RGBADecoderNet, self).__init__()
+        self.conv_rgba = nn.Sequential(nn.Conv2d(c, out_planes, kernel_size=3, stride=1,
+                                                 padding=1, dilation=1, bias=True))
+        self.bottleneck = torch.nn.Sequential()
+        for i in range(num_bottleneck_blocks):
+            self.bottleneck.add_module(
+                'r' + str(i), ResBlock2d(c, kernel_size=(3, 3), padding=(1, 1)))
+
+    def forward(self, features_weighted_mask_atfeaturesscale_list=[]):
+        return torch.sigmoid(self.conv_rgba(self.bottleneck(features_weighted_mask_atfeaturesscale_list.pop(0))))

model/shader.py

@@ -0,0 +1,290 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .warplayer import warp_features
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, in_planes, c=224, out_msgs=0, out_locals=0, block_nums=1, out_masks=1, out_local_flows=32, out_msgs_flows=32, out_feat_flows=0):
+
+        super(DecoderBlock, self).__init__()
+        self.conv0 = nn.Sequential(
+            nn.Conv2d(in_planes, c, 3, 2, 1),
+            nn.PReLU(c),
+            nn.Conv2d(c, c, 3, 2, 1),
+            nn.PReLU(c),
+        )
+
+        self.convblocks = nn.ModuleList()
+        for i in range(block_nums):
+            self.convblocks.append(nn.Sequential(
+                nn.Conv2d(c, c, 3, 1, 1),
+                nn.PReLU(c),
+                nn.Conv2d(c, c, 3, 1, 1),
+                nn.PReLU(c),
+                nn.Conv2d(c, c, 3, 1, 1),
+                nn.PReLU(c),
+                nn.Conv2d(c, c, 3, 1, 1),
+                nn.PReLU(c),
+                nn.Conv2d(c, c, 3, 1, 1),
+                nn.PReLU(c),
+                nn.Conv2d(c, c, 3, 1, 1),
+                nn.PReLU(c),
+            ))
+        self.out_flows = 2
+        self.out_msgs = out_msgs
+        self.out_msgs_flows = out_msgs_flows if out_msgs > 0 else 0
+        self.out_locals = out_locals
+        self.out_local_flows = out_local_flows if out_locals > 0 else 0
+        self.out_masks = out_masks
+        self.out_feat_flows = out_feat_flows
+
+        self.conv_last = nn.Sequential(
+            nn.ConvTranspose2d(c, c, 4, 2, 1),
+            nn.PReLU(c),
+            nn.ConvTranspose2d(c, self.out_flows+self.out_msgs+self.out_msgs_flows +
+                               self.out_locals+self.out_local_flows+self.out_masks+self.out_feat_flows, 4, 2, 1),
+        )
+
+    def forward(self, accumulated_flow, *other):
+        x = [accumulated_flow]
+        for each in other:
+            if each is not None:
+                assert(accumulated_flow.shape[-1] == each.shape[-1]), "decoder want {}, but get {}".format(
+                    accumulated_flow.shape, each.shape)
+                x.append(each)
+        feat = self.conv0(torch.cat(x, dim=1))
+        for convblock1 in self.convblocks:
+            feat = convblock1(feat) + feat
+        feat = self.conv_last(feat)
+        prev = 0
+        flow = feat[:, prev:prev+self.out_flows, :, :]
+        prev += self.out_flows
+        message = feat[:, prev:prev+self.out_msgs,
+                       :, :] if self.out_msgs > 0 else None
+        prev += self.out_msgs
+        message_flow = feat[:, prev:prev + self.out_msgs_flows,
+                            :, :] if self.out_msgs_flows > 0 else None
+        prev += self.out_msgs_flows
+        local_message = feat[:, prev:prev + self.out_locals,
+                             :, :] if self.out_locals > 0 else None
+        prev += self.out_locals
+        local_message_flow = feat[:, prev:prev+self.out_local_flows,
+                                  :, :] if self.out_local_flows > 0 else None
+        prev += self.out_local_flows
+        mask = torch.sigmoid(
+            feat[:, prev:prev+self.out_masks, :, :]) if self.out_masks > 0 else None
+        prev += self.out_masks
+        feat_flow = feat[:, prev:prev+self.out_feat_flows,
+                         :, :] if self.out_feat_flows > 0 else None
+        prev += self.out_feat_flows
+        return flow, mask, message, message_flow, local_message, local_message_flow, feat_flow
+
+
+class CINN(nn.Module):
+    def __init__(self, DIM_SHADER_REFERENCE, target_feature_chns=[512, 256, 128, 64, 64], feature_chns=[2048, 1024, 512, 256, 64], out_msgs_chn=[2048, 1024, 512, 256, 64, 64], out_locals_chn=[2048, 1024, 512, 256, 64, 0], block_num=[1, 1, 1, 1, 1, 2], block_chn_num=[224, 224, 224, 224, 224, 224]):
+        super(CINN, self).__init__()
+       
+        self.in_msgs_chn = [0, *out_msgs_chn[:-1]]
+        self.in_locals_chn = [0, *out_locals_chn[:-1]]
+
+        self.decoder_blocks = nn.ModuleList()
+        self.feed_weighted = True
+        if self.feed_weighted:
+            in_planes = 2+2+DIM_SHADER_REFERENCE*2
+        else:
+            in_planes = 2+DIM_SHADER_REFERENCE
+        for each_target_feature_chns, each_feature_chns, each_out_msgs_chn, each_out_locals_chn, each_in_msgs_chn, each_in_locals_chn, each_block_num, each_block_chn_num in zip(target_feature_chns, feature_chns, out_msgs_chn, out_locals_chn, self.in_msgs_chn, self.in_locals_chn, block_num, block_chn_num):
+            self.decoder_blocks.append(
+                DecoderBlock(in_planes+each_target_feature_chns+each_feature_chns+each_in_locals_chn+each_in_msgs_chn, c=each_block_chn_num, block_nums=each_block_num, out_msgs=each_out_msgs_chn, out_locals=each_out_locals_chn, out_masks=2+each_out_locals_chn))
+        for i in range(len(feature_chns), len(out_locals_chn)):
+            #print("append extra block", i, "msg",
+            #      out_msgs_chn[i], "local", out_locals_chn[i], "block", block_num[i])
+            self.decoder_blocks.append(
+                DecoderBlock(in_planes+self.in_msgs_chn[i]+self.in_locals_chn[i], c=block_chn_num[i], block_nums=block_num[i], out_msgs=out_msgs_chn[i], out_locals=out_locals_chn[i], out_masks=2+out_msgs_chn[i], out_feat_flows=0))
+
+    def apply_flow(self, mask, message, message_flow, local_message, local_message_flow, x_reference, accumulated_flow, each_x_reference_features=None, each_x_reference_features_flow=None):
+        if each_x_reference_features is not None:
+            size_from = each_x_reference_features
+        else:
+            size_from = x_reference
+        f_size = (size_from.shape[2], size_from.shape[3])
+        accumulated_flow = self.flow_rescale(
+            accumulated_flow, size_from)
+        # mask = warp_features(F.interpolate(
+        #    mask, size=f_size, mode="bilinear"), accumulated_flow) if mask is not None else None
+        mask = F.interpolate(
+            mask, size=f_size, mode="bilinear") if mask is not None else None
+        message = F.interpolate(
+            message, size=f_size, mode="bilinear") if message is not None else None
+        message_flow = self.flow_rescale(
+            message_flow, size_from) if message_flow is not None else None
+        message = warp_features(
+            message, message_flow) if message_flow is not None else message
+
+        local_message = F.interpolate(
+            local_message, size=f_size, mode="bilinear") if local_message is not None else None
+        local_message_flow = self.flow_rescale(
+            local_message_flow, size_from) if local_message_flow is not None else None
+        local_message = warp_features(
+            local_message, local_message_flow) if local_message_flow is not None else local_message
+
+        warp_x_reference = warp_features(F.interpolate(
+            x_reference, size=f_size, mode="bilinear"), accumulated_flow)
+
+        each_x_reference_features_flow = self.flow_rescale(
+            each_x_reference_features_flow, size_from) if (each_x_reference_features is not None and each_x_reference_features_flow is not None) else None
+        warp_each_x_reference_features = warp_features(
+            each_x_reference_features, each_x_reference_features_flow) if each_x_reference_features_flow is not None else each_x_reference_features
+
+        return mask, message, local_message, warp_x_reference, accumulated_flow, warp_each_x_reference_features, each_x_reference_features_flow
+
+    def forward(self, x_target_features=[], x_reference=None, x_reference_features=[]):
+        y_flow = []
+        y_feat_flow = []
+
+        y_local_message = []
+        y_warp_x_reference = []
+        y_warp_x_reference_features = []
+
+        y_weighted_flow = []
+        y_weighted_mask = []
+        y_weighted_message = []
+        y_weighted_x_reference = []
+        y_weighted_x_reference_features = []
+
+        for pyrlevel, ifblock in enumerate(self.decoder_blocks):
+            stacked_wref = []
+            stacked_feat = []
+            stacked_anci = []
+            stacked_flow = []
+            stacked_mask = []
+            stacked_mesg = []
+            stacked_locm = []
+            stacked_feat_flow = []
+            for view_id in range(x_reference.shape[1]):  # NMCHW
+
+                if pyrlevel == 0:
+                    # create from zero flow
+                    feat_ev = x_reference_features[pyrlevel][:,
+                                                             view_id, :, :, :] if pyrlevel < len(x_reference_features) else None
+
+                    accumulated_flow = torch.zeros_like(
+                        feat_ev[:, :2, :, :]).to(device)
+                    accumulated_feat_flow = torch.zeros_like(
+                        feat_ev[:, :32, :, :]).to(device)
+                    # domestic inputs
+                    warp_x_reference = F.interpolate(x_reference[:, view_id, :, :, :], size=(
+                        feat_ev.shape[-2], feat_ev.shape[-1]), mode="bilinear")
+                    warp_x_reference_features = feat_ev
+
+                    local_message = None
+                    # federated inputs
+                    weighted_flow = accumulated_flow if self.feed_weighted else None
+                    weighted_wref = warp_x_reference if self.feed_weighted else None
+                    weighted_message = None
+                else:
+                    # resume from last layer
+                    accumulated_flow = y_flow[-1][:, view_id, :, :, :]
+                    accumulated_feat_flow = y_feat_flow[-1][:,
+                                                            view_id, :, :, :] if y_feat_flow[-1] is not None else None
+                    # domestic inputs
+                    warp_x_reference = y_warp_x_reference[-1][:,
+                                                              view_id, :, :, :]
+                    warp_x_reference_features = y_warp_x_reference_features[-1][:,
+                                                                                view_id, :, :, :] if y_warp_x_reference_features[-1] is not None else None
+                    local_message = y_local_message[-1][:, view_id, :,
+                                                        :, :] if len(y_local_message) > 0 else None
+
+                    # federated inputs
+                    weighted_flow = y_weighted_flow[-1] if self.feed_weighted else None
+                    weighted_wref = y_weighted_x_reference[-1] if self.feed_weighted else None
+                    weighted_message = y_weighted_message[-1] if len(
+                        y_weighted_message) > 0 else None
+                scaled_x_target = x_target_features[pyrlevel][:, :, :, :].detach() if pyrlevel < len(
+                    x_target_features) else None
+                # compute flow
+                residual_flow, mask, message, message_flow, local_message, local_message_flow, residual_feat_flow = ifblock(
+                    accumulated_flow, scaled_x_target, warp_x_reference, warp_x_reference_features, weighted_flow, weighted_wref, weighted_message, local_message)
+                accumulated_flow = residual_flow + accumulated_flow
+                accumulated_feat_flow = accumulated_flow
+
+                feat_ev = x_reference_features[pyrlevel+1][:,
+                                                           view_id, :, :, :] if pyrlevel+1 < len(x_reference_features) else None
+                mask, message, local_message, warp_x_reference, accumulated_flow,  warp_x_reference_features, accumulated_feat_flow = self.apply_flow(
+                    mask, message, message_flow, local_message, local_message_flow, x_reference[:, view_id, :, :, :], accumulated_flow, feat_ev, accumulated_feat_flow)
+                stacked_flow.append(accumulated_flow)
+                if accumulated_feat_flow is not None:
+                    stacked_feat_flow.append(accumulated_feat_flow)
+                stacked_mask.append(mask)
+                if message is not None:
+                    stacked_mesg.append(message)
+                if local_message is not None:
+                    stacked_locm.append(local_message)
+                stacked_wref.append(warp_x_reference)
+                if warp_x_reference_features is not None:
+                    stacked_feat.append(warp_x_reference_features)
+
+            stacked_flow = torch.stack(stacked_flow, dim=1)  # M*NCHW -> NMCHW
+            stacked_feat_flow = torch.stack(stacked_feat_flow, dim=1) if len(
+                stacked_feat_flow) > 0 else None
+            stacked_mask = torch.stack(
+                stacked_mask, dim=1)
+            
+            stacked_mesg = torch.stack(stacked_mesg, dim=1) if len(
+                stacked_mesg) > 0 else None
+            stacked_locm = torch.stack(stacked_locm, dim=1) if len(
+                stacked_locm) > 0 else None
+
+            stacked_wref = torch.stack(stacked_wref, dim=1)
+            stacked_feat = torch.stack(stacked_feat, dim=1) if len(
+                stacked_feat) > 0 else None
+            stacked_anci = torch.stack(stacked_anci, dim=1) if len(
+                stacked_anci) > 0 else None
+            y_flow.append(stacked_flow)
+            y_feat_flow.append(stacked_feat_flow)
+
+            y_warp_x_reference.append(stacked_wref)
+            y_warp_x_reference_features.append(stacked_feat)
+            # compute normalized confidence
+            stacked_contrib = torch.nn.functional.softmax(stacked_mask, dim=1)
+
+            # torch.sum to remove temp dimension M from NMCHW --> NCHW
+            weighted_flow = torch.sum(
+                stacked_mask[:, :, 0:1, :, :] * stacked_contrib[:, :, 0:1, :, :] * stacked_flow, dim=1)
+            weighted_mask = torch.sum(
+                stacked_contrib[:, :, 0:1, :, :] * stacked_mask[:, :, 0:1, :, :], dim=1)
+            weighted_wref = torch.sum(
+                stacked_mask[:, :, 0:1, :, :] * stacked_contrib[:, :, 0:1, :, :] * stacked_wref, dim=1) if stacked_wref is not None else None
+            weighted_feat = torch.sum(
+                stacked_mask[:, :, 1:2, :, :] * stacked_contrib[:, :, 1:2, :, :] * stacked_feat, dim=1) if stacked_feat is not None else None
+            weighted_mesg = torch.sum(
+                stacked_mask[:, :, 2:, :, :] * stacked_contrib[:, :, 2:, :, :] * stacked_mesg, dim=1) if stacked_mesg is not None else None
+            y_weighted_flow.append(weighted_flow)
+            y_weighted_mask.append(weighted_mask)
+            if weighted_mesg is not None:
+                y_weighted_message.append(weighted_mesg)
+            if stacked_locm is not None:
+                y_local_message.append(stacked_locm)
+            y_weighted_message.append(weighted_mesg)
+            y_weighted_x_reference.append(weighted_wref)
+            y_weighted_x_reference_features.append(weighted_feat)
+
+            if weighted_feat is not None:
+                y_weighted_x_reference_features.append(weighted_feat)
+        return {
+            "y_last_remote_features": [weighted_mesg],
+        }
+
+    def flow_rescale(self, prev_flow, each_x_reference_features):
+        if prev_flow is None:
+            prev_flow = torch.zeros_like(
+                each_x_reference_features[:, :2]).to(device)
+        else:
+            up_scale_factor = each_x_reference_features.shape[-1] / \
+                prev_flow.shape[-1]
+            if up_scale_factor != 1:
+                prev_flow = F.interpolate(prev_flow, scale_factor=up_scale_factor, mode="bilinear",
+                                          align_corners=False, recompute_scale_factor=False) * up_scale_factor
+        return prev_flow

model/warplayer.py

@@ -0,0 +1,56 @@
+import torch
+import numpy as np
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+backwarp_tenGrid = {}
+
+
+def warp(tenInput, tenFlow):
+    with torch.cuda.amp.autocast(enabled=False):
+        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)
+        if tenInput.dtype != g.dtype:
+            g = g.to(tenInput.dtype)
+        return torch.nn.functional.grid_sample(input=tenInput, grid=g, mode='bilinear', padding_mode='border', align_corners=True)
+# "zeros" "border"
+
+
+def warp_features(inp, flow, ):
+    groups = flow.shape[1]//2  # NCHW
+    samples = inp.shape[0]
+    h = inp.shape[2]
+    w = inp.shape[3]
+    assert(flow.shape[0] == samples and flow.shape[2]
+           == h and flow.shape[3] == w)
+    chns = inp.shape[1]
+    chns_per_group = chns // groups
+    assert(flow.shape[1] % 2 == 0)
+    assert(chns % groups == 0)
+    inp = inp.contiguous().view(samples*groups, chns_per_group, h, w)
+    flow = flow.contiguous().view(samples*groups, 2,  h, w)
+    feat = warp(inp, flow)
+    feat = feat.view(samples, chns, h, w)
+    return feat
+
+
+def flow2rgb(flow_map_np):
+    h, w, _ = flow_map_np.shape
+    rgb_map = np.ones((h, w, 3)).astype(np.float32)/2.0
+    normalized_flow_map = np.concatenate(
+        (flow_map_np[:, :, 0:1]/h/2.0, flow_map_np[:, :, 1:2]/w/2.0), axis=2)
+    rgb_map[:, :, 0] += normalized_flow_map[:, :, 0]
+    rgb_map[:, :, 1] -= 0.5 * \
+        (normalized_flow_map[:, :, 0] + normalized_flow_map[:, :, 1])
+    rgb_map[:, :, 2] += normalized_flow_map[:, :, 1]
+    return (rgb_map.clip(0, 1)*255.0).astype(np.uint8)

poses_template.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d6a3ed9527dd6b04275bf4c81704ac50c9b8cd1bf4a26b1d3ec6658e200ff57
+size 10662874

requirements.txt

@@ -0,0 +1,9 @@
+pillow>=7.2.0
+numpy>=1.16
+tqdm>=4.35.0
+torch>=1.3.0
+opencv-python>=4.5.2
+scikit-image>=0.14.0
+torchvision>=0.2.1
+lpips>=0.1.3
+gdown

train.py

@@ -0,0 +1,203 @@
+import argparse
+import os
+import time
+from datetime import datetime
+from distutils.util import strtobool
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+from torchvision import transforms
+from data_loader import (FileDataset,
+                         RandomResizedCropWithAutoCenteringAndZeroPadding)
+from torch.utils.data.distributed import DistributedSampler
+from conr import CoNR
+from tqdm import tqdm
+
+def data_sampler(dataset, shuffle, distributed):
+
+    if distributed:
+        return torch.utils.data.distributed.DistributedSampler(dataset, shuffle=shuffle)
+
+    if shuffle:
+        return torch.utils.data.RandomSampler(dataset)
+
+    else:
+        return torch.utils.data.SequentialSampler(dataset)
+
+def save_output(image_name, inputs_v, d_dir=".", crop=None):
+    import cv2
+
+    inputs_v = inputs_v.detach().squeeze()
+    input_np = torch.clamp(inputs_v*255, 0, 255).byte().cpu().numpy().transpose(
+        (1, 2, 0))
+    # cv2.setNumThreads(1)
+    out_render_scale = cv2.cvtColor(input_np, cv2.COLOR_RGBA2BGRA)
+    if crop is not None:
+        crop = crop.cpu().numpy()[0]
+        output_img = np.zeros((crop[0], crop[1], 4), dtype=np.uint8)
+        before_resize_scale = cv2.resize(
+            out_render_scale, (crop[5]-crop[4]+crop[8]+crop[9], crop[3]-crop[2]+crop[6]+crop[7]), interpolation=cv2.INTER_AREA)  # w,h
+        output_img[crop[2]:crop[3], crop[4]:crop[5]] = before_resize_scale[crop[6]:before_resize_scale.shape[0] -
+                                                                           crop[7], crop[8]:before_resize_scale.shape[1]-crop[9]]
+    else:
+        output_img = out_render_scale
+    cv2.imwrite(d_dir+"/"+image_name.split(os.sep)[-1]+'.png',
+                output_img
+                )
+
+
+def test():
+    source_names_list = []
+    for name in sorted(os.listdir(args.test_input_person_images)):
+        thissource = os.path.join(args.test_input_person_images, name)
+        if os.path.isfile(thissource):
+            source_names_list.append(thissource)
+        if os.path.isdir(thissource):
+            print("skipping empty folder :"+thissource)
+
+    image_names_list = []
+    for name in sorted(os.listdir(args.test_input_poses_images)):
+        thistarget = os.path.join(args.test_input_poses_images, name)
+        if os.path.isfile(thistarget):
+            image_names_list.append([thistarget, *source_names_list])
+        if os.path.isdir(thistarget):
+            print("skipping folder :"+thistarget)
+    print(image_names_list)
+
+    print("---building models")
+    conrmodel = CoNR(args)
+    conrmodel.load_model(path=args.test_checkpoint_dir)
+    conrmodel.dist()
+    infer(args, conrmodel, image_names_list)
+
+
+def infer(args, humanflowmodel, image_names_list):
+    print("---test images: ", len(image_names_list))
+    test_salobj_dataset = FileDataset(image_names_list=image_names_list,
+                                      fg_img_lbl_transform=transforms.Compose([
+                                          RandomResizedCropWithAutoCenteringAndZeroPadding(
+                                              (args.dataloader_imgsize, args.dataloader_imgsize), scale=(1, 1), ratio=(1.0, 1.0), center_jitter=(0.0, 0.0)
+                                          )]),
+                                      shader_pose_use_gt_udp_test=not args.test_pose_use_parser_udp,
+                                      shader_target_use_gt_rgb_debug=False
+                                      )
+    sampler = data_sampler(test_salobj_dataset, shuffle=False,
+                           distributed=args.distributed)
+    train_data = DataLoader(test_salobj_dataset,
+                            batch_size=1,
+                            shuffle=False,sampler=sampler, 
+                            num_workers=args.dataloaders)
+
+    # start testing
+
+    train_num = train_data.__len__()
+    time_stamp = time.time()
+    prev_frame_rgb = []
+    prev_frame_a = []
+    
+    pbar = tqdm(range(train_num), ncols=100)
+    for i, data in enumerate(train_data):
+        data_time_interval = time.time() - time_stamp
+        time_stamp = time.time()
+        with torch.no_grad():
+            data["character_images"] = torch.cat(
+                [data["character_images"], *prev_frame_rgb], dim=1)
+            data["character_masks"] = torch.cat(
+                [data["character_masks"], *prev_frame_a], dim=1)
+            data = humanflowmodel.data_norm_image(data)
+            pred = humanflowmodel.model_step(data, training=False)
+            # remember to call  humanflowmodel.reset_charactersheet() if you change character .
+
+        train_time_interval = time.time() - time_stamp
+        time_stamp = time.time()
+        if args.local_rank == 0:
+            pbar.set_description(f"Epoch {i}/{train_num}")
+            pbar.set_postfix({"data_time": data_time_interval, "train_time":train_time_interval})
+            pbar.update(1)
+
+        with torch.no_grad():
+
+            if args.test_output_video:
+                pred_img = pred["shader"]["y_weighted_warp_decoded_rgba"]
+                save_output(
+                    str(int(data["imidx"].cpu().item())), pred_img, args.test_output_dir, crop=data["pose_crop"])
+            
+            if args.test_output_udp:
+                pred_img = pred["shader"]["x_target_sudp_a"]
+                save_output(
+                    "udp_"+str(int(data["imidx"].cpu().item())), pred_img, args.test_output_dir)
+
+
+def build_args():
+    parser = argparse.ArgumentParser()
+    # distributed learning settings
+    parser.add_argument("--world_size", type=int, default=1,
+                        help='world size')
+    parser.add_argument("--local_rank", type=int, default=0,
+                        help='local_rank, DON\'T change it')
+
+    # model settings
+    parser.add_argument('--dataloader_imgsize', type=int, default=256,
+                        help='Input image size of the model')
+    parser.add_argument('--batch_size', type=int, default=4,
+                        help='minibatch size')
+    parser.add_argument('--model_name', default='model_result',
+                        help='Name of the experiment')
+    parser.add_argument('--dataloaders', type=int, default=2,
+                        help='Num of dataloaders')
+    parser.add_argument('--mode', default="test", choices=['train', 'test'],
+                        help='Training mode or Testing mode')
+
+    # i/o settings
+    parser.add_argument('--test_input_person_images',
+                        type=str, default="./character_sheet/",
+                        help='Directory to input character sheets')
+    parser.add_argument('--test_input_poses_images', type=str,
+                        default="./test_data/",
+                        help='Directory to input UDP sequences or pose images')
+    parser.add_argument('--test_checkpoint_dir', type=str,
+                        default='./weights/',
+                        help='Directory to model weights')
+    parser.add_argument('--test_output_dir', type=str,
+                        default="./results/",
+                        help='Directory to output images')
+
+    # output content settings
+    parser.add_argument('--test_output_video', type=strtobool, default=True,
+                        help='Whether to output the final result of CoNR, \
+                              images will be output to test_output_dir while True.')
+    parser.add_argument('--test_output_udp', type=strtobool, default=False,
+                        help='Whether to output UDP generated from UDP detector, \
+                              this is meaningful ONLY when test_input_poses_images \
+                              is not UDP sequences but pose images. Meanwhile, \
+                              test_pose_use_parser_udp need to be True')
+
+    # UDP detector settings
+    parser.add_argument('--test_pose_use_parser_udp',
+                        type=strtobool, default=False, 
+                        help='Whether to use UDP detector to generate UDP from pngs, \
+                              pose input MUST be pose images instead of UDP sequences \
+                              while True')
+
+    args = parser.parse_args()
+
+    args.distributed = (args.world_size > 1)
+    if args.local_rank == 0:
+        print("batch_size:", args.batch_size, flush=True)
+    if args.distributed:
+        if args.local_rank == 0:
+            print("world_size: ", args.world_size)
+        torch.distributed.init_process_group(
+            backend="nccl", init_method="env://", world_size=args.world_size)
+        torch.cuda.set_device(args.local_rank)
+        torch.backends.cudnn.benchmark = True
+    else:
+        args.local_rank = 0
+ 
+    return args
+
+
+if __name__ == "__main__":
+    args = build_args()
+    test()

weights/rgbadecodernet.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d1cc48cf4b3b6f7c4856bec8299e57466836eff3bffa73f518965fdb75bc16fd
+size 341897

weights/shader.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5c8fa74e07db0e0a853fe2aa8d299c854b5c957037dcf4858dc0ca755bec9a95
+size 384615535

weights/target_pose_encoder.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a077437779043b543b015d56ca4f521668645c9ad1cd67ee843aa8a94bf59034
+size 107891883

weights/udpparsernet.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:13bb149d18146277cfb11f90b96ae0ecb25828a1ce57ae47ba04d5468cd3322e
+size 228957615