human seg

README.md

@@ -14,6 +14,13 @@ __Contact__: xuebin[at]ualberta[dot]ca
 
 ## Updates !!!
 
+**(2021-Feb-06)** Since some people raised the issue of human segmentation using U^2-Net, we trained another model for human segemntation based on [**Supervisely Person Dataset**](https://supervise.ly/explore/projects/supervisely-person-dataset-23304/datasets). To run the human segmentation model, please first downlowd the [**u2net_human_seg.pth**](https://drive.google.com/file/d/1-Yg0cxgrNhHP-016FPdp902BR-kSsA4P/view?usp=sharing) model weights and start the inference by command: ```python u2net_human_seg_test.py```. <br/>
+[**Notes: Due to the labeling accuracy of the Supervisely Person Dataset, the human segmentation model here won't give you hair level accuracy. But it should be more robust than u2net trained with DUTS-TR dataset on human segmentation task. It can be used for human portrait segmentation, human body segmentation, etc.**](https://github.com/NathanUA/U-2-Net)<br/>
+
+![Human Image Segmentation](figures/human_seg)
+![Human Video](figures/human_seg_video.gif)
+![Human Video Results](figures/human_seg_results.gif)
+
 **(2020-Dec-28)** Some interesting applications and useful tools based on U^2-Net: <br/>
 (1) [**Xiaolong Liu**](https://github.com/LiuXiaolong19920720) developed several very interesting applications based on U^2-Net including [**Human Portrait Drawing**](https://www.cvpy.net/studio/cv/func/DeepLearning/sketch/sketch/page/)(As far as I know, Xiaolong is the first one who uses U^2-Net for portrait generation), [**image matting**](https://www.cvpy.net/studio/cv/func/DeepLearning/matting/matting/page/) and [**so on**](https://www.cvpy.net/). <br/>
 (2) [**Vladimir Seregin**](https://github.com/peko/nn-lineart) developed an interesting tool, [**NN based lineart**](https://peko.github.io/nn-lineart/), for comparing the portrait results of U^2-Net and that of another popular model, [**ArtLine**](https://github.com/vijishmadhavan/ArtLine), developed by [**Vijish Madhavan**](https://github.com/vijishmadhavan). <br/>

u2net_human_seg_test.py

@@ -0,0 +1,117 @@
+import os
+from skimage import io, transform
+import torch
+import torchvision
+from torch.autograd import Variable
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms#, utils
+# import torch.optim as optim
+
+import numpy as np
+from PIL import Image
+import glob
+
+from data_loader import RescaleT
+from data_loader import ToTensor
+from data_loader import ToTensorLab
+from data_loader import SalObjDataset
+
+from model import U2NET # full size version 173.6 MB
+
+# normalize the predicted SOD probability map
+def normPRED(d):
+    ma = torch.max(d)
+    mi = torch.min(d)
+
+    dn = (d-mi)/(ma-mi)
+
+    return dn
+
+def save_output(image_name,pred,d_dir):
+
+    predict = pred
+    predict = predict.squeeze()
+    predict_np = predict.cpu().data.numpy()
+
+    im = Image.fromarray(predict_np*255).convert('RGB')
+    img_name = image_name.split(os.sep)[-1]
+    image = io.imread(image_name)
+    imo = im.resize((image.shape[1],image.shape[0]),resample=Image.BILINEAR)
+
+    pb_np = np.array(imo)
+
+    aaa = img_name.split(".")
+    bbb = aaa[0:-1]
+    imidx = bbb[0]
+    for i in range(1,len(bbb)):
+        imidx = imidx + "." + bbb[i]
+
+    imo.save(d_dir+imidx+'.png')
+
+def main():
+
+    # --------- 1. get image path and name ---------
+    model_name='u2net'
+
+
+    image_dir = os.path.join(os.getcwd(), 'test_data', 'test_human_images')#'test_human_images')#'test_portrait_images', 'your_portrait_im')
+    prediction_dir = os.path.join(os.getcwd(), 'test_data', 'test_human_images' + '_results' + os.sep)
+    model_dir = os.path.join(os.getcwd(), 'saved_models', model_name+'_human_seg', model_name + '_human_seg.pth')
+
+    img_name_list = glob.glob(image_dir + os.sep + '*')
+    print(img_name_list)
+
+    # --------- 2. dataloader ---------
+    #1. dataloader
+    test_salobj_dataset = SalObjDataset(img_name_list = img_name_list,
+                                        lbl_name_list = [],
+                                        transform=transforms.Compose([RescaleT(320),
+                                                                      ToTensorLab(flag=0)])
+                                        )
+    test_salobj_dataloader = DataLoader(test_salobj_dataset,
+                                        batch_size=1,
+                                        shuffle=False,
+                                        num_workers=1)
+
+    # --------- 3. model define ---------
+    if(model_name=='u2net'):
+        print("...load U2NET---173.6 MB")
+        net = U2NET(3,1)
+
+    if torch.cuda.is_available():
+        net.load_state_dict(torch.load(model_dir))
+        net.cuda()
+    else:
+        net.load_state_dict(torch.load(model_dir, map_location='cpu'))
+    net.eval()
+
+    # --------- 4. inference for each image ---------
+    for i_test, data_test in enumerate(test_salobj_dataloader):
+
+        print("inferencing:",img_name_list[i_test].split(os.sep)[-1])
+
+        inputs_test = data_test['image']
+        inputs_test = inputs_test.type(torch.FloatTensor)
+
+        if torch.cuda.is_available():
+            inputs_test = Variable(inputs_test.cuda())
+        else:
+            inputs_test = Variable(inputs_test)
+
+        d1,d2,d3,d4,d5,d6,d7= net(inputs_test)
+
+        # normalization
+        pred = d1[:,0,:,:]
+        pred = normPRED(pred)
+
+        # save results to test_results folder
+        if not os.path.exists(prediction_dir):
+            os.makedirs(prediction_dir, exist_ok=True)
+        save_output(img_name_list[i_test],pred,prediction_dir)
+
+        del d1,d2,d3,d4,d5,d6,d7
+
+if __name__ == "__main__":
+    main()