+# 4. 训练模型
+python src/trainer.py
-- Online Application (在线应用) | - Research Demo | - AAAI 2022 Paper | - Supplementary Video -
- -- Community | - Code | - PPM Benchmark | - License | - Acknowledgement | - Citation | - Contact -
- ---- - - -## 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. - - - -### 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. - - - - -## 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 - [@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), - [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`. - -
diff --git a/pretrained/mobilenetv2_human_seg.ckpt b/pretrained/mobilenetv2_human_seg.ckpt
new file mode 100644
index 0000000..96f2da8
Binary files /dev/null and b/pretrained/mobilenetv2_human_seg.ckpt differ
diff --git a/pretrained/modnet_photographic_portrait_matting.ckpt b/pretrained/modnet_photographic_portrait_matting.ckpt
new file mode 100644
index 0000000..ddb4a4b
Binary files /dev/null and b/pretrained/modnet_photographic_portrait_matting.ckpt differ
diff --git a/src/eval.py b/src/eval.py
new file mode 100644
index 0000000..7106051
--- /dev/null
+++ b/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}')
diff --git a/src/infer.py b/src/infer.py
new file mode 100644
index 0000000..116807b
--- /dev/null
+++ b/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')
+
diff --git a/src/matting_dataset.py b/src/matting_dataset.py
new file mode 100644
index 0000000..4e404fb
--- /dev/null
+++ b/src/matting_dataset.py
@@ -0,0 +1,195 @@
+from torch.utils.data import Dataset
+import numpy as np
+import random
+import cv2
+from glob import glob
+import torch
+from torchvision.transforms import functional as F
+from torchvision import transforms
+from PIL import Image
+
+
+class MattingDataset(Dataset):
+ def __init__(self,
+ dataset_root_dir='src/datasets/PPM-100',
+ transform=None):
+ image_path = dataset_root_dir + '/train/fg/*'
+ matte_path = dataset_root_dir + '/train/alpha/*'
+ image_file_name_list = glob(image_path)
+ matte_file_name_list = glob(matte_path)
+
+ self.image_file_name_list = sorted(image_file_name_list)
+ self.matte_file_name_list = sorted(matte_file_name_list)
+ for img, mat in zip(self.image_file_name_list, self.matte_file_name_list):
+ img_name = img.split('/')[-1]
+ mat_name = mat.split('/')[-1]
+ assert img_name == mat_name
+
+ self.transform = transform
+
+ def __len__(self):
+ return len(self.image_file_name_list)
+
+ def __getitem__(self, index):
+ image_file_name = self.image_file_name_list[index]
+ matte_file_name = self.matte_file_name_list[index]
+
+ image = Image.open(image_file_name)
+ matte = Image.open(matte_file_name)
+
+ data = {'image': image, 'gt_matte': matte}
+
+ if self.transform:
+ data = self.transform(data)
+ return data
+
+
+class Rescale(object):
+ def __init__(self, output_size):
+ assert isinstance(output_size, (int, tuple))
+ self.output_size = output_size
+
+ def __call__(self, sample):
+ image, gt_matte = sample['image'], sample['gt_matte']
+
+ new_h, new_w = int(self.output_size), int(self.output_size)
+ new_img = F.resize(image, (new_h, new_w))
+ new_gt_matte = F.resize(gt_matte, (new_h, new_w))
+
+ return {'image': new_img,'gt_matte': new_gt_matte}
+
+
+class GenTrimap(object):
+ def __call__(self, sample):
+ gt_matte = sample['gt_matte']
+ trimap = self.gen_trimap(gt_matte)
+ sample['trimap'] = trimap
+ return sample
+
+ @staticmethod
+ def gen_trimap(matte):
+ """
+ 根据归matte生成归一化的trimap
+ """
+ matte = np.array(matte)
+ k_size = random.choice(range(2, 5))
+ iterations = np.random.randint(5, 15)
+ kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
+ (k_size, k_size)) # cv2.MORPH_RECT, cv2.MORPH_CROSS
+ dilated = cv2.dilate(matte, kernel, iterations=iterations)
+ eroded = cv2.erode(matte, kernel, iterations=iterations)
+
+ trimap = np.zeros(matte.shape)
+ trimap.fill(0.5)
+ trimap[eroded > 254.5] = 1.0
+ trimap[dilated < 0.5] = 0.0
+ trimap = Image.fromarray(trimap)
+ return trimap
+
+
+class ToTensor(object):
+ """Convert ndarrays in sample to Tensors."""
+
+ def __call__(self, sample):
+ image, trimap, gt_matte = sample['image'], sample['trimap'], sample['gt_matte']
+ image = F.pil_to_tensor(image)
+ trimap = F.pil_to_tensor(trimap)
+ gt_matte = F.pil_to_tensor(gt_matte)
+
+
+ return {'image': image,
+ 'trimap': trimap,
+ 'gt_matte': gt_matte}
+
+
+class ConvertImageDtype(object):
+ def __call__(self, sample):
+ image, trimap, gt_matte = sample['image'], sample['trimap'], sample['gt_matte']
+ image = F.convert_image_dtype(image, torch.float)
+ trimap = F.convert_image_dtype(trimap, torch.float)
+ gt_matte = F.convert_image_dtype(gt_matte, torch.float)
+
+ return {'image': image, 'trimap': trimap, 'gt_matte': gt_matte}
+
+
+class Normalize(object):
+ def __init__(self, mean, std, inplace=False):
+ self.mean = mean
+ self.std = std
+ self.inplace = inplace
+
+ def __call__(self, sample):
+ image, trimap, gt_matte = sample['image'], sample['trimap'], sample['gt_matte']
+ image = image.type(torch.FloatTensor)
+ image = F.normalize(image, self.mean, self.std, self.inplace)
+ sample['image'] = image
+ sample['gt_matte'] = sample['gt_matte'] / 255
+ return sample
+
+
+class ToTrainArray(object):
+ def __call__(self, sample):
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+ image = sample['image'].to(device)
+ trimap = sample['trimap'].to(device)
+ gt_matte = sample['gt_matte'].to(device)
+
+ return [image, trimap, gt_matte]
+ # return [sample['image'], sample['trimap'], sample['gt_matte']]
+
+
+if __name__ == '__main__':
+
+ # test MattingDataset.gen_trimap
+ matte = Image.open('src/datasets/PPM-100/train/alpha/6146816_556eaff97f_o.jpg')
+ trimap1 = GenTrimap().gen_trimap(matte)
+ trimap1 = np.array(trimap1) * 255
+ trimap1 = np.uint8(trimap1)
+ trimap1 = Image.fromarray(trimap1)
+ trimap1.save('test_trimap.png')
+
+ # test MattingDataset
+ transform = transforms.Compose([
+ Rescale(512),
+ GenTrimap(),
+ ToTensor(),
+ # Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+ ])
+ mattingDataset = MattingDataset(transform=transform)
+
+ import matplotlib.pyplot as plt
+
+ fig = plt.figure()
+
+ for i in range(len(mattingDataset)):
+ sample = mattingDataset[i]
+ print(mattingDataset.image_file_name_list[i])
+ # print(sample)
+ print(i, sample['image'].shape, sample['trimap'].shape, sample['gt_matte'].shape)
+
+ # break
+
+ ax = plt.subplot(4, 3, 3 * i + 1)
+ plt.tight_layout()
+ ax.set_title('image #{}'.format(i))
+ ax.axis('off')
+ img = transforms.ToPILImage()(sample['image'])
+ plt.imshow(img)
+
+ ax = plt.subplot(4, 3, 3 * i + 2)
+ plt.tight_layout()
+ ax.set_title('gt_matte #{}'.format(i))
+ ax.axis('off')
+ img = transforms.ToPILImage()(sample['gt_matte'])
+ plt.imshow(img)
+
+ ax = plt.subplot(4, 3, 3 * i + 3)
+ plt.tight_layout()
+ ax.set_title('trimap #{}'.format(i))
+ ax.axis('off')
+ img = transforms.ToPILImage()(sample['trimap'])
+ plt.imshow(img)
+
+ if i == 3:
+ plt.show()
+ break
diff --git a/src/requirements.txt b/src/requirements.txt
new file mode 100644
index 0000000..d1b497d
--- /dev/null
+++ b/src/requirements.txt
@@ -0,0 +1,7 @@
+torch==1.10.2
+scipy==1.7.3
+numpy==1.21.5
+opencv-python==4.5.5.62
+matplotlib==3.5.1
+torchvision==0.11.3
+Pillow==9.0.1
\ No newline at end of file
diff --git a/src/setting.py b/src/setting.py
new file mode 100644
index 0000000..9e2070d
--- /dev/null
+++ b/src/setting.py
@@ -0,0 +1,9 @@
+BS = 16 # BATCH SIZE
+LR = 0.01 # LEARN RATE
+EPOCHS = 4 # TOTAL EPOCHS
+
+SEMANTIC_SCALE = 10.0
+DETAIL_SCALE = 10.0
+MATTE_SCALE = 1.0
+
+SAVE_EPOCH_STEP = 3
diff --git a/src/trainer.py b/src/trainer.py
index bd3d8be..becb4af 100644
--- a/src/trainer.py
+++ b/src/trainer.py
@@ -7,7 +7,6 @@
import torch.nn as nn
import torch.nn.functional as F
-
__all__ = [
'supervised_training_iter',
'soc_adaptation_iter',
@@ -80,7 +79,9 @@ def _init_kernel(self):
# MODNet Training Functions
# ----------------------------------------------------------------------------------
-blurer = GaussianBlurLayer(1, 3).cuda()
+blurer = GaussianBlurLayer(1, 3) #.cuda
+if torch.cuda.is_available():
+ blurer.cuda()
def supervised_training_iter(
@@ -297,3 +298,59 @@ def soc_adaptation_iter(
return soc_semantic_loss, soc_detail_loss
# ----------------------------------------------------------------------------------
+
+
+if __name__ == '__main__':
+ from matting_dataset import MattingDataset, Rescale, \
+ ToTensor, Normalize, ToTrainArray, \
+ ConvertImageDtype, GenTrimap
+ from torchvision import transforms
+ from torch.utils.data import DataLoader
+ from models.modnet import MODNet
+ from setting import BS, LR, EPOCHS, SEMANTIC_SCALE, DETAIL_SCALE, MATTE_SCALE, SAVE_EPOCH_STEP
+ transform = transforms.Compose([
+ Rescale(512),
+ GenTrimap(),
+ ToTensor(),
+ ConvertImageDtype(),
+ Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+ ToTrainArray()
+ ])
+ mattingDataset = MattingDataset(transform=transform)
+
+ modnet = torch.nn.DataParallel(MODNet()) #.cuda()
+ if torch.cuda.is_available():
+ modnet = modnet.cuda()
+
+ optimizer = torch.optim.SGD(modnet.parameters(), lr=LR, momentum=0.9)
+ lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=int(0.25 * EPOCHS), gamma=0.1)
+
+ dataloader = DataLoader(mattingDataset,
+ batch_size=BS,
+ shuffle=True)
+
+ for epoch in range(0, EPOCHS):
+ print(f'epoch: {epoch}/{EPOCHS-1}')
+ for idx, (image, trimap, gt_matte) in enumerate(dataloader):
+ semantic_loss, detail_loss, matte_loss = \
+ supervised_training_iter(modnet, optimizer, image, trimap, gt_matte,
+ semantic_scale=SEMANTIC_SCALE,
+ detail_scale=DETAIL_SCALE,
+ matte_scale=MATTE_SCALE)
+ print(f'{(idx+1) * BS}/{len(mattingDataset)} --- '
+ f'semantic_loss: {semantic_loss:f}, detail_loss: {detail_loss:f}, matte_loss: {matte_loss:f}\r',
+ end='')
+ lr_scheduler.step()
+ # 保存中间训练结果
+ if epoch % SAVE_EPOCH_STEP == 0 and epoch > 1:
+ torch.save({
+ 'epoch': epoch,
+ 'model_state_dict': modnet.state_dict(),
+ 'optimizer_state_dict': optimizer.state_dict(),
+ 'loss': {'semantic_loss': semantic_loss, 'detail_loss': detail_loss, 'matte_loss': matte_loss},
+ }, f'pretrained/modnet_custom_portrait_matting_{epoch}_th.ckpt')
+ print(f'{len(mattingDataset)}/{len(mattingDataset)} --- '
+ f'semantic_loss: {semantic_loss:f}, detail_loss: {detail_loss:f}, matte_loss: {matte_loss:f}')
+
+ # 仅保存模型权重参数
+ torch.save(modnet.state_dict(), f'pretrained/modnet_custom_portrait_matting_last_epoch_weight.ckpt')