[Feature] Support Style-Based Global Appearance Flow for Virtual Try-On

.gitignore

@@ -147,3 +147,7 @@ resources/
 # Slurm
 batchscript-*
 *.out
+
+*.png
+*.jpg
+*.zip

projects/flow_style_vton/.gitignore

@@ -0,0 +1,2 @@
+*.jpg
+*.zip

projects/flow_style_vton/README.md

@@ -0,0 +1,96 @@
+# Style-Based Global Appearance Flow for Virtual Try-On (CVPR 2022)
+
+## Description
+
+Awesome try-on desplays are like this:
+
+![image1](examples/000010_0.png)
+
+```
+Author: @FerryHuang.
+
+This is an implementation of https://github.com/SenHe/Flow-Style-VTON adapting to mmediting. Only inference is supported so far.
+```
+
+## Usage
+
+### Setup Environment
+
+Please refer to [Get Started](https://mmediting.readthedocs.io/en/latest/get_started/I.html) to install
+MMEditing.
+
+At first, add the current folder to `PYTHONPATH`, so that Python can find your code. Run command in the current directory to add it.
+
+> Please run it every time after you opened a new shell.
+
+```shell
+export PYTHONPATH=`pwd`:$PYTHONPATH
+```
+
+### Data Preparation
+
+Please check the [official repo](https://github.com/SenHe/Flow-Style-VTON) and download test-set and pretrained checkpoints and put them under the folder projects/flow_style_vton
+
+### Testing commands
+
+**To test with single GPU:**
+
+```bash
+cd projects/flow_style_vton
+python inference.py
+```
+
+Expectedly, two folders will be made im_gar_flow_wg and our_t_results, containing the
+try-on procedures and the final results, respectively.
+
+## Citation
+
+<!-- Replace to the citation of the paper your project refers to. -->
+
+```bibtex
+@inproceedings{he2022fs_vton,
+  title={Style-Based Global Appearance Flow for Virtual Try-On},
+  author={He, Sen and Song, Yi-Zhe and Xiang, Tao},
+  booktitle={CVPR},
+  year={2022}
+}
+```
+
+## Checklist \[required\]
+
+Here is a checklist of this project's progress. And you can ignore this part if you don't plan to contribute
+to MMediting projects.
+
+- [ ] Milestone 1: PR-ready, and acceptable to be one of the `projects/`.
+
+  - [ ] Finish the code
+
+    <!-- The code's design shall follow existing interfaces and convention. For example, each model component should be registered into `mmedit.registry.MODELS` and configurable via a config file. -->
+
+  - [ ] Basic docstrings & proper citation
+
+    <!-- Each major class should contains a docstring, describing its functionality and arguments. If your code is copied or modified from other open-source projects, don't forget to cite the source project in docstring and make sure your behavior is not against its license. Typically, we do not accept any code snippet under GPL license. [A Short Guide to Open Source Licenses](https://medium.com/nationwide-technology/a-short-guide-to-open-source-licenses-cf5b1c329edd) -->
+
+  - [ ] Converted checkpoint and results (Only for reproduction)
+
+    <!-- If you are reproducing the result from a paper, make sure the model in the project can match that results. Also please provide checkpoint links or a checkpoint conversion script for others to get the pre-trained model. -->
+
+- [ ] Milestone 2: Indicates a successful model implementation.
+
+  - [ ] Training results
+
+    <!-- If you are reproducing the result from a paper, train your model from scratch and verified that the final result can match the original result. -->
+
+- [ ] Milestone 3: Good to be a part of our core package!
+
+  - [ ] Unit tests
+
+    <!-- Unit tests for the major module are required. [Example](https://github.com/open-mmlab/mmediting/blob/main/tests/test_models/test_backbones/test_vision_transformer.py) -->
+
+  - [ ] Code style
+
+    <!-- Refactor your code according to reviewer's comment. -->
+
+  - [ ] `metafile.yml` and `README.md`
+
+    <!-- It will used for mmediting to acquire your models. [Example](https://github.com/open-mmlab/mmediting/blob/main/configs/mvit/metafile.yml). In particular, you may have to refactor this README into a standard one. [Example](https://github.com/open-mmlab/mmediting/blob/main/configs/swin_transformer/README.md) -->

projects/flow_style_vton/configs/flow_style_vton_PFAFN_epoch_101.py

@@ -0,0 +1,8 @@
+model = dict(
+    type='FlowStyleVTON',
+    warp_model=dict(type='AFWM', input_nc=3),
+    gen_model=dict(
+        type='ResUnetGenerator', input_nc=7, output_nc=4, num_downs=5),
+    pretrained_cfgs=dict(
+        warp_model=dict(ckpt_path='ckp/aug/PFAFN_warp_epoch_101.pth'),
+        gen_model=dict(ckpt_path='ckp/aug/PFAFN_gen_epoch_101.pth')))

projects/flow_style_vton/inference.py

@@ -0,0 +1,69 @@
+import os
+from argparse import ArgumentParser
+
+import torch
+from mmengine.registry import init_default_scope
+from torch.utils.data import DataLoader
+from torchvision.utils import save_image
+from tqdm import tqdm
+from vton_dataset import AlignedDataset
+
+from mmagic.apis.inferencers.inference_functions import init_model
+from projects.flow_style_vton.models import FlowStyleVTON
+
+init_default_scope('mmedit')
+
+config = 'configs/flow_style_vton_PFAFN_epoch_101.py'
+
+parser = ArgumentParser()
+parser.add_argument('--gpu_id', type=int, default=0)
+parser.add_argument(
+    '--loadSize', type=int, default=512, help='scale images to this size')
+parser.add_argument(
+    '--fineSize', type=int, default=512, help='then crop to this size')
+parser.add_argument('--dataroot', type=str, default='VITON_test')
+parser.add_argument('--test_pairs', type=str, default='test_pairs.txt')
+parser.add_argument(
+    '--resize_or_crop',
+    type=str,
+    default='scale_width',
+    help='scaling and cropping of images at load time \
+    [resize_and_crop|crop|scale_width|scale_width_and_crop]')
+parser.add_argument('--phase', type=str, default='test')
+parser.add_argument('--isTrain', default=False)
+parser.add_argument(
+    '--no_flip',
+    action='store_true',
+    help='if specified, do not flip the images for data argumentation')
+parser.add_argument('--batch_size', type=int, default=1)
+parser.add_argument('--num_workers', type=int, default=4)
+
+parser.add_argument('--output_dir', type=str, default='inference_results')
+opt = parser.parse_args()
+
+dataset = AlignedDataset(opt)
+dataloader = DataLoader(dataset, opt.batch_size, num_workers=opt.num_workers)
+
+device = torch.device(
+    f'cuda:{opt.gpu_id}' if torch.cuda.is_available() else 'cpu')
+# pretrained is set inside the config
+model = init_model(config).to(device).eval()
+assert isinstance(model, FlowStyleVTON)
+
+os.makedirs('our_t_results', exist_ok=True)
+os.makedirs('im_gar_flow_wg', exist_ok=True)
+for i, data in enumerate(tqdm(dataloader)):
+    p_tryon, combine = model.infer(data)
+    save_image(
+        p_tryon,
+        os.path.join('our_t_results', data['p_name'][0]),
+        nrow=int(1),
+        normalize=True,
+        value_range=(-1, 1))
+    save_image(
+        combine,
+        os.path.join('im_gar_flow_wg', data['p_name'][0]),
+        nrow=int(1),
+        normalize=True,
+        range=(-1, 1),
+    )

projects/flow_style_vton/models/__init__.py

@@ -0,0 +1,5 @@
+from .afwm import AFWM
+from .flow_style_vton_model import FlowStyleVTON
+from .generator import ResUnetGenerator
+
+__all__ = ['FlowStyleVTON', 'AFWM', 'ResUnetGenerator']