问题1

.gitignore

@@ -94,4 +94,6 @@ ENV/
 
 
 # Project files
-.vscode
+.vscode
+.idea/
+src/datasets/PPM-100

README.md

@@ -1,126 +1,24 @@
-<h2 align="center">MODNet: Trimap-Free Portrait Matting in Real Time</h2>
+# 说明
+代码fork from [MODNet官方代码](https://github.com/ZHKKKe/MODNet) 。本项目完善了数据准备、模型评价及模型训练相关代码
+# 模型训练、评价、推理
+```bash
+# 1. 下载代码并进入工作目录
+git clone https://github.com/actboy/MODNet
+cd MODNet
 
-<div align="center"><i>MODNet: Real-Time Trimap-Free Portrait Matting via Objective Decomposition (AAAI 2022)</i></div>
+# 2. 安装依赖
+pip install -r src/requirements.txt
 
-<br />
+# 3. 下载并解压数据集
+wget -c https://paddleseg.bj.bcebos.com/matting/datasets/PPM-100.zip -O src/datasets/PPM-100.zip
+unzip src/datasets/PPM-100.zip -d src/datasets
 
-<img src="doc/gif/homepage_demo.gif" width="100%">
+# 4. 训练模型
+python src/trainer.py
 
-<div align="center">MODNet is a model for <b>real-time</b> portrait matting with <b>only RGB image input</b></div>
-<div align="center">MODNet是一个<b>仅需RGB图片输入</b>的<b>实时</b>人像抠图模型</div>
+# 5. 模型评估
+python src/eval.py
 
-<br />
-
-<p align="center">
-  <a href="#online-application-在线应用">Online Application (在线应用)</a> |
-  <a href="#research-demo">Research Demo</a> | 
-  <a href="https://arxiv.org/pdf/2011.11961.pdf">AAAI 2022 Paper</a> |
-  <a href="https://youtu.be/PqJ3BRHX3Lc">Supplementary Video</a>
-</p>
-
-<p align="center">
-  <a href="#community">Community</a> |
-  <a href="#code">Code</a> |
-  <a href="#ppm-benchmark">PPM Benchmark</a> | 
-  <a href="#license">License</a> | 
-  <a href="#acknowledgement">Acknowledgement</a> | 
-  <a href="#citation">Citation</a> | 
-  <a href="#contact">Contact</a>
-</p>
-
----
-
-
-## Online Application (在线应用)
-
-A **Single** model! Only **7M**! Process **2K** resolution image with a **Fast** speed on common PCs or Mobiles! **Beter** than research demos!  
-Please try online portrait image matting via [this website](https://sight-x.cn/portrait_matting)!    
-
-**单个**模型！大小仅为**7M**！可以在普通PC或移动设备上**快速**处理具有**2K**分辨率的图像！效果比研究示例**更好**！  
-请通过[此网站](https://sight-x.cn/portrait_matting)在线尝试图片抠像！
-
-
-## Research Demo
-
-All the models behind the following demos are trained on the datasets mentioned in [our paper](https://arxiv.org/pdf/2011.11961.pdf). 
-
-### Portrait Image Matting
-We provide an [online Colab demo](https://colab.research.google.com/drive/1GANpbKT06aEFiW-Ssx0DQnnEADcXwQG6?usp=sharing) for portrait image matting.  
-It allows you to upload portrait images and predict/visualize/download the alpha mattes. 
-
-<!-- <img src="doc/gif/image_matting_demo.gif" width='40%'> -->
-
-### Portrait Video Matting
-We provide two real-time portrait video matting demos based on WebCam. When using the demo, you can move the WebCam around at will. 
-If you have an Ubuntu system, we recommend you to try the [offline demo](demo/video_matting/webcam) to get a higher *fps*. Otherwise, you can access the [online Colab demo](https://colab.research.google.com/drive/1Pt3KDSc2q7WxFvekCnCLD8P0gBEbxm6J?usp=sharing).  
-We also provide an [offline demo](demo/video_matting/custom) that allows you to process custom videos.
-
-<!-- <img src="doc/gif/video_matting_demo.gif" width='60%'> -->
-
-
-## Community
-
-We share some cool applications/extentions of MODNet built by the community.
-
-- **WebGUI for Portrait Image Matting**  
-You can try [this WebGUI](https://www.gradio.app/hub/aliabd/modnet) (hosted on [Gradio](https://www.gradio.app/)) for portrait image matting from your browser without code! 
-
-- **Colab Demo of Bokeh (Blur Background)**  
-You can try [this Colab demo](https://colab.research.google.com/github/eyaler/avatars4all/blob/master/yarok.ipynb) (built by [@eyaler](https://github.com/eyaler)) to blur the backgroud based on MODNet!
-
-- **ONNX Version of MODNet**  
-You can convert the pre-trained MODNet to an ONNX model by using [this code](onnx) (provided by [@manthan3C273](https://github.com/manthan3C273)). You can also try [this Colab demo](https://colab.research.google.com/drive/1P3cWtg8fnmu9karZHYDAtmm1vj1rgA-f?usp=sharing) for MODNet image matting (ONNX version).
-
-- **TorchScript Version of MODNet**  
-You can convert the pre-trained MODNet to an TorchScript model by using [this code](torchscript) (provided by [@yarkable](https://github.com/yarkable)). 
-
-- **TensorRT Version of MODNet**  
-You can access [this Github repository](https://github.com/jkjung-avt/tensorrt_demos) to try the TensorRT version of MODNet (provided by [@jkjung-avt](https://github.com/jkjung-avt)).
-
-
-There are some resources about MODNet from the community.
-- [Video from What's AI YouTube Channel](https://youtu.be/rUo0wuVyefU)
-- [Article from Louis Bouchard's Blog](https://www.louisbouchard.ai/remove-background/)
-
-
-## Code
-We provide the [code](src/trainer.py) of MODNet training iteration, including:
-- **Supervised Training**: Train MODNet on a labeled matting dataset
-- **SOC Adaptation**: Adapt a trained MODNet to an unlabeled dataset
-
-In code comments, we provide examples for using the functions.  
-
-
-## PPM Benchmark
-The PPM benchmark is released in a separate repository [PPM](https://github.com/ZHKKKe/PPM).  
-
-
-## License
-The code, models, and demos in this repository (excluding GIF files under the folder `doc/gif`) are released under the [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0) license.  
-
-
-## Acknowledgement  
-- We thank  
-&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;[@yzhou0919](https://github.com/yzhou0919), [@eyaler](https://github.com/eyaler), [@manthan3C273](https://github.com/manthan3C273),  [@yarkable](https://github.com/yarkable), [@jkjung-avt](https://github.com/jkjung-avt),  [@manzke](https://github.com/manzke),  
-&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;[the Gradio team](https://github.com/gradio-app/gradio), [What's AI YouTube Channel](https://www.youtube.com/channel/UCUzGQrN-lyyc0BWTYoJM_Sg), [Louis Bouchard's Blog](https://www.louisbouchard.ai),  
-for their contributions to this repository or their cool applications/extentions/resources of MODNet.
-
-
-## Citation
-If this work helps your research, please consider to cite:
-
-```bibtex
-@InProceedings{MODNet,
-  author = {Zhanghan Ke and Jiayu Sun and Kaican Li and Qiong Yan and Rynson W.H. Lau},
-  title = {MODNet: Real-Time Trimap-Free Portrait Matting via Objective Decomposition},
-  booktitle = {AAAI},
-  year = {2022},
-}
+# 6. 模型推理
+python src/infer.py
 ```
-
-
-## Contact
-This repository is currently maintained by Zhanghan Ke ([@ZHKKKe](https://github.com/ZHKKKe)).  
-For questions, please contact `kezhanghan@outlook.com`.
-
-<img src="doc/gif/commercial_image_matting_model_result.gif" width='100%'>

src/eval.py

@@ -0,0 +1,72 @@
+import numpy as np
+from glob import glob
+from models.modnet import MODNet
+from PIL import Image
+from infer import predit_matte
+import torch.nn as nn
+import torch
+
+
+def cal_mad(pred, gt):
+    diff = pred - gt
+    diff = np.abs(diff)
+    mad = np.mean(diff)
+    return mad
+
+
+def cal_mse(pred, gt):
+    diff = pred - gt
+    diff = diff ** 2
+    mse = np.mean(diff)
+    return mse
+
+
+def load_eval_dataset(dataset_root_dir='src/datasets/PPM-100'):
+    image_path = dataset_root_dir + '/val/fg/*'
+    matte_path = dataset_root_dir + '/val/alpha/*'
+    image_file_name_list = glob(image_path)
+    image_file_name_list = sorted(image_file_name_list)
+    matte_file_name_list = glob(matte_path)
+    matte_file_name_list = sorted(matte_file_name_list)
+
+    return image_file_name_list, matte_file_name_list
+
+
+def eval(modnet: MODNet, dataset):
+    mse = total_mse = 0.0
+    mad = total_mad = 0.0
+    cnt = 0
+
+    for im_pth, mt_pth in zip(dataset[0], dataset[1]):
+        im = Image.open(im_pth)
+        pd_matte = predit_matte(modnet, im)
+
+        gt_matte = Image.open(mt_pth)
+        gt_matte = np.asarray(gt_matte) / 255
+
+        total_mse += cal_mse(pd_matte, gt_matte)
+        total_mad += cal_mad(pd_matte, gt_matte)
+
+        cnt += 1
+    if cnt > 0:
+        mse = total_mse / cnt
+        mad = total_mad / cnt
+
+    return mse, mad
+
+
+if __name__ == '__main__':
+    # create MODNet and load the pre-trained ckpt
+    modnet = MODNet(backbone_pretrained=False)
+    modnet = nn.DataParallel(modnet)
+
+    ckp_pth = 'pretrained/modnet_photographic_portrait_matting.ckpt'
+    if torch.cuda.is_available():
+        modnet = modnet.cuda()
+        weights = torch.load(ckp_pth)
+    else:
+        weights = torch.load(ckp_pth, map_location=torch.device('cpu'))
+    modnet.load_state_dict(weights)
+    dataset = load_eval_dataset('src/datasets/PPM-100')
+    mse, mad = eval(modnet, dataset)
+    print(f'mse: {mse:6f}, mad: {mad:6f}')

src/infer.py

@@ -0,0 +1,85 @@
+from models.modnet import MODNet
+from PIL import Image
+import numpy as np
+from torchvision import transforms
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+
+
+def predit_matte(modnet: MODNet, im: Image):
+    # define image to tensor transform
+    im_transform = transforms.Compose(
+        [
+            transforms.ToTensor(),
+            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+        ]
+    )
+
+    # define hyper-parameters
+    ref_size = 512
+
+    modnet.eval()
+
+    # unify image channels to 3
+    im = np.asarray(im)
+    if len(im.shape) == 2:
+        im = im[:, :, None]
+    if im.shape[2] == 1:
+        im = np.repeat(im, 3, axis=2)
+    elif im.shape[2] == 4:
+        im = im[:, :, 0:3]
+
+    im = Image.fromarray(im)
+    # convert image to PyTorch tensor
+    im = im_transform(im)
+
+    # add mini-batch dim
+    im = im[None, :, :, :]
+
+    # resize image for input
+    im_b, im_c, im_h, im_w = im.shape
+    if max(im_h, im_w) < ref_size or min(im_h, im_w) > ref_size:
+        if im_w >= im_h:
+            im_rh = ref_size
+            im_rw = int(im_w / im_h * ref_size)
+        elif im_w < im_h:
+            im_rw = ref_size
+            im_rh = int(im_h / im_w * ref_size)
+    else:
+        im_rh = im_h
+        im_rw = im_w
+
+    im_rw = im_rw - im_rw % 32
+    im_rh = im_rh - im_rh % 32
+    im = F.interpolate(im, size=(im_rh, im_rw), mode='area')
+
+    # inference
+    _, _, matte = modnet(im.cuda() if torch.cuda.is_available() else im, True)
+
+    # resize and save matte
+    matte = F.interpolate(matte, size=(im_h, im_w), mode='area')
+    matte = matte[0][0].data.cpu().numpy()
+    return matte
+
+
+if __name__ == '__main__':
+    # create MODNet and load the pre-trained ckpt
+    modnet = MODNet(backbone_pretrained=False)
+    modnet = nn.DataParallel(modnet)
+
+    ckp_pth = 'pretrained/modnet_photographic_portrait_matting.ckpt'
+    if torch.cuda.is_available():
+        modnet = modnet.cuda()
+        weights = torch.load(ckp_pth)
+    else:
+        weights = torch.load(ckp_pth, map_location=torch.device('cpu'))
+    modnet.load_state_dict(weights)
+
+    pth = 'src/datasets/PPM-100/val/fg/5588688353_3426d4b5d9_o.jpg'
+    img = Image.open(pth)
+
+    matte = predit_matte(modnet, img)
+    prd_img = Image.fromarray(((matte * 255).astype('uint8')), mode='L')
+    prd_img.save('test_predic.jpg')
+