├── README.md ├── image ├── README.md └── sample.png ├── model.py ├── results ├── joint │ ├── photo-1438761681033-6461ffad8d80.jpg │ ├── photo-1473830394358-91588751b241.jpg │ ├── photo-1491349174775-aaafddd81942.jpg │ ├── photo-1494790108377-be9c29b29330.jpg │ ├── photo-1499952127939-9bbf5af6c51c.jpg │ ├── photo-1500648767791-00dcc994a43e.jpg │ ├── photo-1503023345310-bd7c1de61c7d.jpg │ ├── photo-1503249023995-51b0f3778ccf.jpg │ ├── photo-1542206395-9feb3edaa68d.jpg │ ├── photo-1544005313-94ddf0286df2.jpg │ ├── photo-1544723795-3fb6469f5b39.jpg │ ├── photo-1547425260-76bcadfb4f2c.jpg │ ├── photo-1552058544-f2b08422138a.jpg │ ├── photo-1554151228-14d9def656e4.jpg │ ├── photo-1554727242-741c14fa561c.jpg │ ├── photo-1567532939604-b6b5b0db2604.jpg │ ├── photo-1573140247632-f8fd74997d5c.jpg │ ├── photo-1593104547489-5cfb3839a3b5.jpg │ └── photo-1599566150163-29194dcaad36.jpg └── mask │ ├── photo-1438761681033-6461ffad8d80.jpg │ ├── photo-1473830394358-91588751b241.jpg │ ├── photo-1491349174775-aaafddd81942.jpg │ ├── photo-1494790108377-be9c29b29330.jpg │ ├── photo-1499952127939-9bbf5af6c51c.jpg │ ├── photo-1500648767791-00dcc994a43e.jpg │ ├── photo-1503023345310-bd7c1de61c7d.jpg │ ├── photo-1503249023995-51b0f3778ccf.jpg │ ├── photo-1542206395-9feb3edaa68d.jpg │ ├── photo-1544005313-94ddf0286df2.jpg │ ├── photo-1544723795-3fb6469f5b39.jpg │ ├── photo-1547425260-76bcadfb4f2c.jpg │ ├── photo-1552058544-f2b08422138a.jpg │ ├── photo-1554151228-14d9def656e4.jpg │ ├── photo-1554727242-741c14fa561c.jpg │ ├── photo-1567532939604-b6b5b0db2604.jpg │ ├── photo-1573140247632-f8fd74997d5c.jpg │ ├── photo-1593104547489-5cfb3839a3b5.jpg │ └── photo-1599566150163-29194dcaad36.jpg ├── test.py ├── test ├── photo-1438761681033-6461ffad8d80.jpg ├── photo-1473830394358-91588751b241.jpg ├── photo-1491349174775-aaafddd81942.jpg ├── photo-1494790108377-be9c29b29330.jpg ├── photo-1499952127939-9bbf5af6c51c.jpg ├── photo-1500648767791-00dcc994a43e.jpg ├── photo-1503023345310-bd7c1de61c7d.jpg ├── photo-1503249023995-51b0f3778ccf.jpg ├── photo-1542206395-9feb3edaa68d.jpg ├── photo-1544005313-94ddf0286df2.jpg ├── photo-1544723795-3fb6469f5b39.jpg ├── photo-1547425260-76bcadfb4f2c.jpg ├── photo-1552058544-f2b08422138a.jpg ├── photo-1554151228-14d9def656e4.jpg ├── photo-1554727242-741c14fa561c.jpg ├── photo-1567532939604-b6b5b0db2604.jpg ├── photo-1573140247632-f8fd74997d5c.jpg ├── photo-1593104547489-5cfb3839a3b5.jpg └── photo-1599566150163-29194dcaad36.jpg └── train.py /README.md: -------------------------------------------------------------------------------- 1 | # U2-Net for Image Matting in TensorFlow 2 | 3 | Welcome to the U2-Net for Image Matting in TensorFlow repository! This project showcases an implementation of the U2-Net architecture for image matting in the TensorFlow. The code includes the training process, making use of the Privacy-Preserving Portrait Matting Dataset (P3M-10k) to achieve accurate foreground extraction from images. 4 | 5 | ## Dataset 6 | Privacy-Preserving Portrait Matting Dataset (P3M-10k) is used for training and validation process. P3M-10k contains 10421 high-resolution real-world face-blurred portrait images, along with their manually labeled alpha mattes.

7 | Download the dataset: 8 | - [Privacy-Preserving Portrait Matting Dataset (P3M-10k)](https://drive.google.com/uc?export=download&id=1LqUU7BZeiq8I3i5KxApdOJ2haXm-cEv1) 9 | - [P3M-10k facemask (optional)](https://drive.google.com/file/d/1I-71PbkWcivBv3ly60V0zvtYRd3ddyYs/view?usp=sharing) 10 | 11 | 12 | 13 | ## Results 14 | The sequence in the images below is as follows- `Input Image`, `Predicted Alpha Matte` and `Predicted Alpha Matte applied over Input Image`.

15 | ![](results/joint/photo-1494790108377-be9c29b29330.jpg) 16 | ![](results/joint/photo-1544723795-3fb6469f5b39.jpg) 17 | ![](results/joint/photo-1438761681033-6461ffad8d80.jpg) 18 | 19 | ## How to improve 20 | - Train on more epochs. 21 | - Increase the input image resolution. 22 | - Apply data augmentation. 23 | - Try new loss function. 24 | 25 | 26 | ## Contact 27 | For more follow me on: 28 | 29 | - YouTube 30 | - Facebook 31 | - Twitter 32 | - Instagram 33 | - Telegram 34 | -------------------------------------------------------------------------------- /image/README.md: -------------------------------------------------------------------------------- 1 | # Images 2 | -------------------------------------------------------------------------------- /image/sample.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/image/sample.png -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2" 4 | 5 | import tensorflow as tf 6 | from tensorflow.keras.layers import Input, Conv2D, BatchNormalization, Activation, MaxPool2D, UpSampling2D, Concatenate, Add 7 | 8 | def conv_block(inputs, out_ch, rate=1): 9 | x = Conv2D(out_ch, 3, padding="same", dilation_rate=1)(inputs) 10 | x = BatchNormalization()(x) 11 | x = Activation("relu")(x) 12 | return x 13 | 14 | def RSU_L(inputs, out_ch, int_ch, num_layers, rate=2): 15 | """ Initial Conv """ 16 | x = conv_block(inputs, out_ch) 17 | init_feats = x 18 | 19 | """ Encoder """ 20 | skip = [] 21 | x = conv_block(x, int_ch) 22 | skip.append(x) 23 | 24 | for i in range(num_layers-2): 25 | x = MaxPool2D((2, 2))(x) 26 | x = conv_block(x, int_ch) 27 | skip.append(x) 28 | 29 | """ Bridge """ 30 | x = conv_block(x, int_ch, rate=rate) 31 | 32 | """ Decoder """ 33 | skip.reverse() 34 | 35 | x = Concatenate()([x, skip[0]]) 36 | x = conv_block(x, int_ch) 37 | 38 | for i in range(num_layers-3): 39 | x = UpSampling2D(size=(2, 2), interpolation="bilinear")(x) 40 | x = Concatenate()([x, skip[i+1]]) 41 | x = conv_block(x, int_ch) 42 | 43 | x = UpSampling2D(size=(2, 2), interpolation="bilinear")(x) 44 | x = Concatenate()([x, skip[-1]]) 45 | x = conv_block(x, out_ch) 46 | 47 | """ Add """ 48 | x = Add()([x, init_feats]) 49 | return x 50 | 51 | def RSU_4F(inputs, out_ch, int_ch): 52 | """ Initial Conv """ 53 | x0 = conv_block(inputs, out_ch, rate=1) 54 | 55 | """ Encoder """ 56 | x1 = conv_block(x0, int_ch, rate=1) 57 | x2 = conv_block(x1, int_ch, rate=2) 58 | x3 = conv_block(x2, int_ch, rate=4) 59 | 60 | """ Bridge """ 61 | x4 = conv_block(x3, int_ch, rate=8) 62 | 63 | """ Decoder """ 64 | x = Concatenate()([x4, x3]) 65 | x = conv_block(x, int_ch, rate=4) 66 | 67 | x = Concatenate()([x, x2]) 68 | x = conv_block(x, int_ch, rate=2) 69 | 70 | x = Concatenate()([x, x1]) 71 | x = conv_block(x, out_ch, rate=1) 72 | 73 | """ Addition """ 74 | x = Add()([x, x0]) 75 | return x 76 | 77 | def u2net(input_shape, out_ch, int_ch, num_classes=1): 78 | """ Input Layer """ 79 | inputs = Input(input_shape) 80 | s0 = inputs 81 | 82 | """ Encoder """ 83 | s1 = RSU_L(s0, out_ch[0], int_ch[0], 7) 84 | p1 = MaxPool2D((2, 2))(s1) 85 | 86 | s2 = RSU_L(p1, out_ch[1], int_ch[1], 6) 87 | p2 = MaxPool2D((2, 2))(s2) 88 | 89 | s3 = RSU_L(p2, out_ch[2], int_ch[2], 5) 90 | p3 = MaxPool2D((2, 2))(s3) 91 | 92 | s4 = RSU_L(p3, out_ch[3], int_ch[3], 4) 93 | p4 = MaxPool2D((2, 2))(s4) 94 | 95 | s5 = RSU_4F(p4, out_ch[4], int_ch[4]) 96 | p5 = MaxPool2D((2, 2))(s5) 97 | 98 | """ Bridge """ 99 | b1 = RSU_4F(p5, out_ch[5], int_ch[5]) 100 | b2 = UpSampling2D(size=(2, 2), interpolation="bilinear")(b1) 101 | 102 | """ Decoder """ 103 | d1 = Concatenate()([b2, s5]) 104 | d1 = RSU_4F(d1, out_ch[6], int_ch[6]) 105 | u1 = UpSampling2D(size=(2, 2), interpolation="bilinear")(d1) 106 | 107 | d2 = Concatenate()([u1, s4]) 108 | d2 = RSU_L(d2, out_ch[7], int_ch[7], 4) 109 | u2 = UpSampling2D(size=(2, 2), interpolation="bilinear")(d2) 110 | 111 | d3 = Concatenate()([u2, s3]) 112 | d3 = RSU_L(d3, out_ch[8], int_ch[8], 5) 113 | u3 = UpSampling2D(size=(2, 2), interpolation="bilinear")(d3) 114 | 115 | d4 = Concatenate()([u3, s2]) 116 | d4 = RSU_L(d4, out_ch[9], int_ch[9], 6) 117 | u4 = UpSampling2D(size=(2, 2), interpolation="bilinear")(d4) 118 | 119 | d5 = Concatenate()([u4, s1]) 120 | d5 = RSU_L(d5, out_ch[10], int_ch[10], 7) 121 | 122 | """ Side Outputs """ 123 | y1 = Conv2D(num_classes, 3, padding="same")(d5) 124 | 125 | y2 = Conv2D(num_classes, 3, padding="same")(d4) 126 | y2 = UpSampling2D(size=(2, 2), interpolation="bilinear")(y2) 127 | 128 | y3 = Conv2D(num_classes, 3, padding="same")(d3) 129 | y3 = UpSampling2D(size=(4, 4), interpolation="bilinear")(y3) 130 | 131 | y4 = Conv2D(num_classes, 3, padding="same")(d2) 132 | y4 = UpSampling2D(size=(8, 8), interpolation="bilinear")(y4) 133 | 134 | y5 = Conv2D(num_classes, 3, padding="same")(d1) 135 | y5 = UpSampling2D(size=(16, 16), interpolation="bilinear")(y5) 136 | 137 | y6 = Conv2D(num_classes, 3, padding="same")(b1) 138 | y6 = UpSampling2D(size=(32, 32), interpolation="bilinear")(y6) 139 | 140 | y0 = Concatenate()([y1, y2, y3, y4, y5, y6]) 141 | y0 = Conv2D(num_classes, 3, padding="same")(y0) 142 | 143 | y0 = Activation("sigmoid", name="y0")(y0) 144 | y1 = Activation("sigmoid", name="y1")(y1) 145 | y2 = Activation("sigmoid", name="y2")(y2) 146 | y3 = Activation("sigmoid", name="y3")(y3) 147 | y4 = Activation("sigmoid", name="y4")(y4) 148 | y5 = Activation("sigmoid", name="y5")(y5) 149 | y6 = Activation("sigmoid", name="y6")(y6) 150 | 151 | model = tf.keras.models.Model(inputs, outputs=[y0, y1, y2, y3, y4, y5, y6]) 152 | return model 153 | 154 | def build_u2net(input_shape, num_classes=1): 155 | out_ch = [64, 128, 256, 512, 512, 512, 512, 256, 128, 64, 64] 156 | int_ch = [32, 32, 64, 128, 256, 256, 256, 128, 64, 32, 16] 157 | model = u2net(input_shape, out_ch, int_ch, num_classes=num_classes) 158 | return model 159 | 160 | def build_u2net_lite(input_shape, num_classes=1): 161 | out_ch = [64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64] 162 | int_ch = [16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16] 163 | model = u2net(input_shape, out_ch, int_ch, num_classes=num_classes) 164 | return model 165 | 166 | if __name__ == "__main__": 167 | model = build_u2net_lite((512, 512, 3)) 168 | model.summary() 169 | 170 | 171 | 172 | 173 | 174 | 175 | 176 | ## ... 177 | -------------------------------------------------------------------------------- /results/joint/photo-1438761681033-6461ffad8d80.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1438761681033-6461ffad8d80.jpg -------------------------------------------------------------------------------- /results/joint/photo-1473830394358-91588751b241.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1473830394358-91588751b241.jpg -------------------------------------------------------------------------------- /results/joint/photo-1491349174775-aaafddd81942.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1491349174775-aaafddd81942.jpg -------------------------------------------------------------------------------- /results/joint/photo-1494790108377-be9c29b29330.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1494790108377-be9c29b29330.jpg -------------------------------------------------------------------------------- /results/joint/photo-1499952127939-9bbf5af6c51c.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1499952127939-9bbf5af6c51c.jpg -------------------------------------------------------------------------------- /results/joint/photo-1500648767791-00dcc994a43e.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1500648767791-00dcc994a43e.jpg -------------------------------------------------------------------------------- /results/joint/photo-1503023345310-bd7c1de61c7d.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1503023345310-bd7c1de61c7d.jpg -------------------------------------------------------------------------------- /results/joint/photo-1503249023995-51b0f3778ccf.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1503249023995-51b0f3778ccf.jpg -------------------------------------------------------------------------------- /results/joint/photo-1542206395-9feb3edaa68d.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1542206395-9feb3edaa68d.jpg -------------------------------------------------------------------------------- /results/joint/photo-1544005313-94ddf0286df2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1544005313-94ddf0286df2.jpg -------------------------------------------------------------------------------- /results/joint/photo-1544723795-3fb6469f5b39.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1544723795-3fb6469f5b39.jpg -------------------------------------------------------------------------------- /results/joint/photo-1547425260-76bcadfb4f2c.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1547425260-76bcadfb4f2c.jpg -------------------------------------------------------------------------------- /results/joint/photo-1552058544-f2b08422138a.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1552058544-f2b08422138a.jpg -------------------------------------------------------------------------------- /results/joint/photo-1554151228-14d9def656e4.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1554151228-14d9def656e4.jpg -------------------------------------------------------------------------------- /results/joint/photo-1554727242-741c14fa561c.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1554727242-741c14fa561c.jpg -------------------------------------------------------------------------------- /results/joint/photo-1567532939604-b6b5b0db2604.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1567532939604-b6b5b0db2604.jpg -------------------------------------------------------------------------------- /results/joint/photo-1573140247632-f8fd74997d5c.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1573140247632-f8fd74997d5c.jpg -------------------------------------------------------------------------------- /results/joint/photo-1593104547489-5cfb3839a3b5.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1593104547489-5cfb3839a3b5.jpg -------------------------------------------------------------------------------- /results/joint/photo-1599566150163-29194dcaad36.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/joint/photo-1599566150163-29194dcaad36.jpg -------------------------------------------------------------------------------- /results/mask/photo-1438761681033-6461ffad8d80.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1438761681033-6461ffad8d80.jpg -------------------------------------------------------------------------------- /results/mask/photo-1473830394358-91588751b241.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1473830394358-91588751b241.jpg -------------------------------------------------------------------------------- /results/mask/photo-1491349174775-aaafddd81942.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1491349174775-aaafddd81942.jpg -------------------------------------------------------------------------------- /results/mask/photo-1494790108377-be9c29b29330.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1494790108377-be9c29b29330.jpg -------------------------------------------------------------------------------- /results/mask/photo-1499952127939-9bbf5af6c51c.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1499952127939-9bbf5af6c51c.jpg -------------------------------------------------------------------------------- /results/mask/photo-1500648767791-00dcc994a43e.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1500648767791-00dcc994a43e.jpg -------------------------------------------------------------------------------- /results/mask/photo-1503023345310-bd7c1de61c7d.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1503023345310-bd7c1de61c7d.jpg -------------------------------------------------------------------------------- /results/mask/photo-1503249023995-51b0f3778ccf.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1503249023995-51b0f3778ccf.jpg -------------------------------------------------------------------------------- /results/mask/photo-1542206395-9feb3edaa68d.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1542206395-9feb3edaa68d.jpg -------------------------------------------------------------------------------- /results/mask/photo-1544005313-94ddf0286df2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1544005313-94ddf0286df2.jpg -------------------------------------------------------------------------------- /results/mask/photo-1544723795-3fb6469f5b39.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1544723795-3fb6469f5b39.jpg -------------------------------------------------------------------------------- /results/mask/photo-1547425260-76bcadfb4f2c.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1547425260-76bcadfb4f2c.jpg -------------------------------------------------------------------------------- /results/mask/photo-1552058544-f2b08422138a.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1552058544-f2b08422138a.jpg -------------------------------------------------------------------------------- /results/mask/photo-1554151228-14d9def656e4.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1554151228-14d9def656e4.jpg -------------------------------------------------------------------------------- /results/mask/photo-1554727242-741c14fa561c.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1554727242-741c14fa561c.jpg -------------------------------------------------------------------------------- /results/mask/photo-1567532939604-b6b5b0db2604.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1567532939604-b6b5b0db2604.jpg -------------------------------------------------------------------------------- /results/mask/photo-1573140247632-f8fd74997d5c.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1573140247632-f8fd74997d5c.jpg -------------------------------------------------------------------------------- /results/mask/photo-1593104547489-5cfb3839a3b5.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1593104547489-5cfb3839a3b5.jpg -------------------------------------------------------------------------------- /results/mask/photo-1599566150163-29194dcaad36.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/results/mask/photo-1599566150163-29194dcaad36.jpg -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2" 4 | 5 | import numpy as np 6 | import cv2 7 | import pandas as pd 8 | from glob import glob 9 | from tqdm import tqdm 10 | import tensorflow as tf 11 | from train import load_dataset 12 | 13 | """ Global parameters """ 14 | H = 512 15 | W = 512 16 | 17 | """ Creating a directory """ 18 | def create_dir(path): 19 | if not os.path.exists(path): 20 | os.makedirs(path) 21 | 22 | if __name__ == "__main__": 23 | """ Seeding """ 24 | np.random.seed(42) 25 | tf.random.set_seed(42) 26 | 27 | """ Directory for storing files """ 28 | for item in ["joint", "mask"]: 29 | create_dir(f"results/{item}") 30 | 31 | """ Load the model """ 32 | model_path = os.path.join("files", "model.h5") 33 | model = tf.keras.models.load_model(model_path) 34 | 35 | """ Dataset """ 36 | images = glob("test/*") 37 | # images = glob("/media/nikhil/New Volume/ML_DATASET/people_segmentation/images/*") 38 | print(f"Images: {len(images)}") 39 | 40 | """ Prediction """ 41 | for x in tqdm(images, total=len(images)): 42 | """ Extracting the name """ 43 | name = x.split("/")[-1] 44 | 45 | """ Reading the image """ 46 | image = cv2.imread(x, cv2.IMREAD_COLOR) 47 | x = cv2.resize(image, (W, H)) 48 | x = x/255.0 49 | x = np.expand_dims(x, axis=0) 50 | 51 | """ Prediction """ 52 | pred = model.predict(x, verbose=0) 53 | 54 | line = np.ones((H, 10, 3)) * 255 55 | 56 | """ Joint and save mask """ 57 | pred_list = [] 58 | for item in pred: 59 | p = item[0] * 255 60 | p = np.concatenate([p, p, p], axis=-1) 61 | 62 | pred_list.append(p) 63 | pred_list.append(line) 64 | 65 | save_image_path = os.path.join("results", "mask", name) 66 | cat_images = np.concatenate(pred_list, axis=1) 67 | cv2.imwrite(save_image_path, cat_images) 68 | 69 | """ Save final mask """ 70 | image_h, image_w, _ = image.shape 71 | 72 | y0 = pred[0][0] 73 | y0 = cv2.resize(y0, (image_w, image_h)) 74 | y0 = np.expand_dims(y0, axis=-1) 75 | y0 = np.concatenate([y0, y0, y0], axis=-1) 76 | 77 | line = line = np.ones((image_h, 10, 3)) * 255 78 | 79 | cat_images = np.concatenate([image, line, y0*255, line, image*y0], axis=1) 80 | save_image_path = os.path.join("results", "joint", name) 81 | cv2.imwrite(save_image_path, cat_images) 82 | -------------------------------------------------------------------------------- /test/photo-1438761681033-6461ffad8d80.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1438761681033-6461ffad8d80.jpg -------------------------------------------------------------------------------- /test/photo-1473830394358-91588751b241.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1473830394358-91588751b241.jpg -------------------------------------------------------------------------------- /test/photo-1491349174775-aaafddd81942.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1491349174775-aaafddd81942.jpg -------------------------------------------------------------------------------- /test/photo-1494790108377-be9c29b29330.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1494790108377-be9c29b29330.jpg -------------------------------------------------------------------------------- /test/photo-1499952127939-9bbf5af6c51c.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1499952127939-9bbf5af6c51c.jpg -------------------------------------------------------------------------------- /test/photo-1500648767791-00dcc994a43e.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1500648767791-00dcc994a43e.jpg -------------------------------------------------------------------------------- /test/photo-1503023345310-bd7c1de61c7d.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1503023345310-bd7c1de61c7d.jpg -------------------------------------------------------------------------------- /test/photo-1503249023995-51b0f3778ccf.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1503249023995-51b0f3778ccf.jpg -------------------------------------------------------------------------------- /test/photo-1542206395-9feb3edaa68d.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1542206395-9feb3edaa68d.jpg -------------------------------------------------------------------------------- /test/photo-1544005313-94ddf0286df2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1544005313-94ddf0286df2.jpg -------------------------------------------------------------------------------- /test/photo-1544723795-3fb6469f5b39.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1544723795-3fb6469f5b39.jpg -------------------------------------------------------------------------------- /test/photo-1547425260-76bcadfb4f2c.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1547425260-76bcadfb4f2c.jpg -------------------------------------------------------------------------------- /test/photo-1552058544-f2b08422138a.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1552058544-f2b08422138a.jpg -------------------------------------------------------------------------------- /test/photo-1554151228-14d9def656e4.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1554151228-14d9def656e4.jpg -------------------------------------------------------------------------------- /test/photo-1554727242-741c14fa561c.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1554727242-741c14fa561c.jpg -------------------------------------------------------------------------------- /test/photo-1567532939604-b6b5b0db2604.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1567532939604-b6b5b0db2604.jpg -------------------------------------------------------------------------------- /test/photo-1573140247632-f8fd74997d5c.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1573140247632-f8fd74997d5c.jpg -------------------------------------------------------------------------------- /test/photo-1593104547489-5cfb3839a3b5.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1593104547489-5cfb3839a3b5.jpg -------------------------------------------------------------------------------- /test/photo-1599566150163-29194dcaad36.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/nikhilroxtomar/U2-Net-for-Image-Matting-in-TensorFlow/c708f054f0d147f7ba478898b0a8e8eb7561939e/test/photo-1599566150163-29194dcaad36.jpg -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2" 4 | 5 | import numpy as np 6 | import cv2 7 | from glob import glob 8 | from sklearn.utils import shuffle 9 | import tensorflow as tf 10 | from tensorflow.keras.callbacks import ModelCheckpoint, CSVLogger, ReduceLROnPlateau, EarlyStopping, TensorBoard 11 | from tensorflow.keras.optimizers import Adam 12 | from sklearn.model_selection import train_test_split 13 | from model import build_u2net_lite, build_u2net 14 | 15 | """ Global parameters """ 16 | H = 512 17 | W = 512 18 | 19 | def create_dir(path): 20 | if not os.path.exists(path): 21 | os.makedirs(path) 22 | 23 | def load_dataset(path, split=0.1): 24 | train_x = sorted(glob(os.path.join(path, "train", "blurred_image", "*.jpg"))) 25 | train_y = sorted(glob(os.path.join(path, "train", "mask", "*.png"))) 26 | 27 | valid_x = sorted(glob(os.path.join(path, "validation", "P3M-500-NP", "original_image", "*.jpg"))) 28 | valid_y = sorted(glob(os.path.join(path, "validation", "P3M-500-NP", "mask", "*.png"))) 29 | 30 | return (train_x, train_y), (valid_x, valid_y) 31 | 32 | def read_image(path): 33 | path = path.decode() 34 | x = cv2.imread(path, cv2.IMREAD_COLOR) 35 | x = cv2.resize(x, (W, H)) 36 | x = x / 255.0 37 | x = x.astype(np.float32) 38 | return x 39 | 40 | def read_mask(path): 41 | path = path.decode() 42 | x = cv2.imread(path, cv2.IMREAD_GRAYSCALE) 43 | x = cv2.resize(x, (W, H)) 44 | x = x / 255.0 45 | x = x.astype(np.float32) 46 | x = np.expand_dims(x, axis=-1) 47 | return x 48 | 49 | def tf_parse(x, y): 50 | def _parse(x, y): 51 | x = read_image(x) 52 | y = read_mask(y) 53 | return x, y 54 | 55 | x, y = tf.numpy_function(_parse, [x, y], [tf.float32, tf.float32]) 56 | x.set_shape([H, W, 3]) 57 | y.set_shape([H, W, 1]) 58 | return x, y 59 | 60 | def tf_dataset(X, Y, batch=2): 61 | ds = tf.data.Dataset.from_tensor_slices((X, Y)) 62 | ds = ds.map(tf_parse).batch(batch).prefetch(10) 63 | return ds 64 | 65 | if __name__ == "__main__": 66 | """ Seeding """ 67 | np.random.seed(42) 68 | tf.random.set_seed(42) 69 | 70 | """ Directory for storing files """ 71 | create_dir("files") 72 | 73 | """ Hyperparameters """ 74 | batch_size = 8 75 | lr = 1e-4 76 | num_epochs = 500 77 | model_path = os.path.join("files", "model.h5") 78 | csv_path = os.path.join("files", "log.csv") 79 | 80 | """ Dataset """ 81 | dataset_path = "/media/nikhil/New Volume/ML_DATASET/IMAGE-MATTING/P3M-10k" 82 | (train_x, train_y), (valid_x, valid_y) = load_dataset(dataset_path) 83 | 84 | print(f"Train: {len(train_x)} - {len(train_y)}") 85 | print(f"Valid: {len(valid_x)} - {len(valid_y)}") 86 | 87 | train_dataset = tf_dataset(train_x, train_y, batch=batch_size) 88 | valid_dataset = tf_dataset(valid_x, valid_y, batch=batch_size) 89 | 90 | """ Model """ 91 | model = build_u2net((H, W, 3)) 92 | model.load_weights(model_path) 93 | model.compile(loss="binary_crossentropy", optimizer=Adam(lr)) 94 | 95 | callbacks = [ 96 | ModelCheckpoint(model_path, verbose=1, save_best_only=True), 97 | ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=5, min_lr=1e-7, verbose=1), 98 | CSVLogger(csv_path), 99 | EarlyStopping(monitor='val_loss', patience=20, restore_best_weights=False), 100 | ] 101 | 102 | model.fit( 103 | train_dataset, 104 | epochs=num_epochs, 105 | validation_data=valid_dataset, 106 | callbacks=callbacks 107 | ) 108 | --------------------------------------------------------------------------------