wip: init commit

head_segmentation/face_parsing/.gitignore

@@ -0,0 +1,5 @@
+*.mp4
+*.jpg
+*.png
+*.zip
+*.avi

head_segmentation/face_parsing/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 zll
+
+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.

head_segmentation/face_parsing/README.md

@@ -0,0 +1,68 @@
+# face-parsing.PyTorch
+
+<p align="center">
+	<a href="https://github.com/zllrunning/face-parsing.PyTorch">
+    <img class="page-image" src="https://github.com/zllrunning/face-parsing.PyTorch/blob/master/6.jpg" >
+	</a>
+</p>
+
+### Contents
+- [Training](#training)
+- [Demo](#Demo)
+- [References](#references)
+
+## Training
+
+1. Prepare training data:
+    -- download [CelebAMask-HQ dataset](https://github.com/switchablenorms/CelebAMask-HQ)
+
+	--  change file path in the `prepropess_data.py`  and run
+```Shell
+python prepropess_data.py
+```
+
+2. Train the model using CelebAMask-HQ dataset:
+Just run the train script: 
+```
+    $ CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py
+```
+
+If you do not wish to train the model, you can download [our pre-trained model](https://drive.google.com/open?id=154JgKpzCPW82qINcVieuPH3fZ2e0P812) and save it in `res/cp`.
+
+
+## Demo
+1. Evaluate the trained model using:
+```Shell
+# evaluate using GPU
+python test.py
+```
+
+## Face makeup using parsing maps
+[**face-makeup.PyTorch**](https://github.com/zllrunning/face-makeup.PyTorch)
+<table>
+
+<tr>
+<th>&nbsp;</th>
+<th>Hair</th>
+<th>Lip</th>
+</tr>
+
+<!-- Line 1: Original Input -->
+<tr>
+<td><em>Original Input</em></td>
+<td><img src="makeup/116_ori.png" height="256" width="256" alt="Original Input"></td>
+<td><img src="makeup/116_lip_ori.png" height="256" width="256" alt="Original Input"></td>
+</tr>
+
+<!-- Line 3: Color -->
+<tr>
+<td>Color</td>
+<td><img src="makeup/116_1.png" height="256" width="256" alt="Color"></td>
+<td><img src="makeup/116_3.png" height="256" width="256" alt="Color"></td>
+</tr>
+
+</table>
+
+
+## References
+- [BiSeNet](https://github.com/CoinCheung/BiSeNet)

head_segmentation/face_parsing/evaluate.py

@@ -0,0 +1,95 @@
+#!/usr/bin/python
+# -*- encoding: utf-8 -*-
+
+from logger import setup_logger
+from model import BiSeNet
+from face_dataset import FaceMask
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+import torch.distributed as dist
+
+import os
+import os.path as osp
+import logging
+import time
+import numpy as np
+from tqdm import tqdm
+import math
+from PIL import Image
+import torchvision.transforms as transforms
+import cv2
+
+def vis_parsing_maps(im, parsing_anno, stride, save_im=False, save_path='vis_results/parsing_map_on_im.jpg'):
+    # Colors for all 20 parts
+    part_colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0],
+                   [255, 0, 85], [255, 0, 170],
+                   [0, 255, 0], [85, 255, 0], [170, 255, 0],
+                   [0, 255, 85], [0, 255, 170],
+                   [0, 0, 255], [85, 0, 255], [170, 0, 255],
+                   [0, 85, 255], [0, 170, 255],
+                   [255, 255, 0], [255, 255, 85], [255, 255, 170],
+                   [255, 0, 255], [255, 85, 255], [255, 170, 255],
+                   [0, 255, 255], [85, 255, 255], [170, 255, 255]]
+
+    im = np.array(im)
+    vis_im = im.copy().astype(np.uint8)
+    vis_parsing_anno = parsing_anno.copy().astype(np.uint8)
+    vis_parsing_anno = cv2.resize(vis_parsing_anno, None, fx=stride, fy=stride, interpolation=cv2.INTER_NEAREST)
+    vis_parsing_anno_color = np.zeros((vis_parsing_anno.shape[0], vis_parsing_anno.shape[1], 3)) + 255
+
+    num_of_class = np.max(vis_parsing_anno)
+
+    for pi in range(1, num_of_class + 1):
+        index = np.where(vis_parsing_anno == pi)
+        vis_parsing_anno_color[index[0], index[1], :] = part_colors[pi]
+
+    vis_parsing_anno_color = vis_parsing_anno_color.astype(np.uint8)
+    # print(vis_parsing_anno_color.shape, vis_im.shape)
+    vis_im = cv2.addWeighted(cv2.cvtColor(vis_im, cv2.COLOR_RGB2BGR), 0.4, vis_parsing_anno_color, 0.6, 0)
+
+    # Save result or not
+    if save_im:
+        cv2.imwrite(save_path, vis_im, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
+
+    # return vis_im
+
+def evaluate(respth='./res/test_res', dspth='./data', cp='model_final_diss.pth'):
+
+    if not os.path.exists(respth):
+        os.makedirs(respth)
+
+    n_classes = 19
+    net = BiSeNet(n_classes=n_classes)
+    net.cuda()
+    save_pth = osp.join('res/cp', cp)
+    net.load_state_dict(torch.load(save_pth))
+    net.eval()
+
+    to_tensor = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    ])
+    with torch.no_grad():
+        for image_path in os.listdir(dspth):
+            img = Image.open(osp.join(dspth, image_path))
+            image = img.resize((512, 512), Image.BILINEAR)
+            img = to_tensor(image)
+            img = torch.unsqueeze(img, 0)
+            img = img.cuda()
+            out = net(img)[0]
+            parsing = out.squeeze(0).cpu().numpy().argmax(0)
+
+            vis_parsing_maps(image, parsing, stride=1, save_im=True, save_path=osp.join(respth, image_path))
+
+
+
+
+
+
+
+if __name__ == "__main__":
+    setup_logger('./res')
+    evaluate()

head_segmentation/face_parsing/face_dataset.py

@@ -0,0 +1,106 @@
+#!/usr/bin/python
+# -*- encoding: utf-8 -*-
+
+import torch
+from torch.utils.data import Dataset
+import torchvision.transforms as transforms
+
+import os.path as osp
+import os
+from PIL import Image
+import numpy as np
+import json
+import cv2
+
+from transform import *
+
+
+
+class FaceMask(Dataset):
+    def __init__(self, rootpth, cropsize=(640, 480), mode='train', *args, **kwargs):
+        super(FaceMask, self).__init__(*args, **kwargs)
+        assert mode in ('train', 'val', 'test')
+        self.mode = mode
+        self.ignore_lb = 255
+        self.rootpth = rootpth
+
+        self.imgs = os.listdir(os.path.join(self.rootpth, 'CelebA-HQ-img'))
+
+        #  pre-processing
+        self.to_tensor = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+            ])
+        self.trans_train = Compose([
+            ColorJitter(
+                brightness=0.5,
+                contrast=0.5,
+                saturation=0.5),
+            HorizontalFlip(),
+            RandomScale((0.75, 1.0, 1.25, 1.5, 1.75, 2.0)),
+            RandomCrop(cropsize)
+            ])
+
+    def __getitem__(self, idx):
+        impth = self.imgs[idx]
+        img = Image.open(osp.join(self.rootpth, 'CelebA-HQ-img', impth))
+        img = img.resize((512, 512), Image.BILINEAR)
+        label = Image.open(osp.join(self.rootpth, 'mask', impth[:-3]+'png')).convert('P')
+        # print(np.unique(np.array(label)))
+        if self.mode == 'train':
+            im_lb = dict(im=img, lb=label)
+            im_lb = self.trans_train(im_lb)
+            img, label = im_lb['im'], im_lb['lb']
+        img = self.to_tensor(img)
+        label = np.array(label).astype(np.int64)[np.newaxis, :]
+        return img, label
+
+    def __len__(self):
+        return len(self.imgs)
+
+
+if __name__ == "__main__":
+    face_data = '/home/zll/data/CelebAMask-HQ/CelebA-HQ-img'
+    face_sep_mask = '/home/zll/data/CelebAMask-HQ/CelebAMask-HQ-mask-anno'
+    mask_path = '/home/zll/data/CelebAMask-HQ/mask'
+    counter = 0
+    total = 0
+    for i in range(15):
+        # files = os.listdir(osp.join(face_sep_mask, str(i)))
+
+        atts = ['skin', 'l_brow', 'r_brow', 'l_eye', 'r_eye', 'eye_g', 'l_ear', 'r_ear', 'ear_r',
+                'nose', 'mouth', 'u_lip', 'l_lip', 'neck', 'neck_l', 'cloth', 'hair', 'hat']
+
+        for j in range(i*2000, (i+1)*2000):
+
+            mask = np.zeros((512, 512))
+
+            for l, att in enumerate(atts, 1):
+                total += 1
+                file_name = ''.join([str(j).rjust(5, '0'), '_', att, '.png'])
+                path = osp.join(face_sep_mask, str(i), file_name)
+
+                if os.path.exists(path):
+                    counter += 1
+                    sep_mask = np.array(Image.open(path).convert('P'))
+                    # print(np.unique(sep_mask))
+
+                    mask[sep_mask == 225] = l
+            cv2.imwrite('{}/{}.png'.format(mask_path, j), mask)
+            print(j)
+
+    print(counter, total)
+
+
+
+
+
+
+
+
+
+
+
+
+
+

head_segmentation/face_parsing/logger.py

@@ -0,0 +1,23 @@
+#!/usr/bin/python
+# -*- encoding: utf-8 -*-
+
+
+import os.path as osp
+import time
+import sys
+import logging
+
+import torch.distributed as dist
+
+
+def setup_logger(logpth):
+    logfile = 'BiSeNet-{}.log'.format(time.strftime('%Y-%m-%d-%H-%M-%S'))
+    logfile = osp.join(logpth, logfile)
+    FORMAT = '%(levelname)s %(filename)s(%(lineno)d): %(message)s'
+    log_level = logging.INFO
+    if dist.is_initialized() and not dist.get_rank()==0:
+        log_level = logging.ERROR
+    logging.basicConfig(level=log_level, format=FORMAT, filename=logfile)
+    logging.root.addHandler(logging.StreamHandler())
+
+