├── Assets ├── U2NETPRmodel.png ├── U_2_NETP_IMAGE3_background_changed.png ├── U_2_NETP_IMAGE3_cropped_no-bg.jpg ├── U_2_NETP_IMAGE3_cropped_no-bg_mask.jpg ├── U_2_NETP_Results1.jpg ├── U_2_NETP_Results2.jpg ├── U_2_NETP_Results3.jpg └── u2netexamples.png ├── DEMOS ├── U2NET_COMPLETE.ipynb ├── U_2_Netp_Cropper_Colab.ipynb ├── U_2_Netp_Demonstration_Colab.ipynb └── U_2_Netp_Step_by_Step_Demonstration_Colab_v2.ipynb ├── LICENSE └── README.md /Assets/U2NETPRmodel.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shreyas-bk/u2netdemo/3958d19a822115c6e39520e38e0303c11f299d44/Assets/U2NETPRmodel.png -------------------------------------------------------------------------------- /Assets/U_2_NETP_IMAGE3_background_changed.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shreyas-bk/u2netdemo/3958d19a822115c6e39520e38e0303c11f299d44/Assets/U_2_NETP_IMAGE3_background_changed.png -------------------------------------------------------------------------------- /Assets/U_2_NETP_IMAGE3_cropped_no-bg.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shreyas-bk/u2netdemo/3958d19a822115c6e39520e38e0303c11f299d44/Assets/U_2_NETP_IMAGE3_cropped_no-bg.jpg -------------------------------------------------------------------------------- /Assets/U_2_NETP_IMAGE3_cropped_no-bg_mask.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shreyas-bk/u2netdemo/3958d19a822115c6e39520e38e0303c11f299d44/Assets/U_2_NETP_IMAGE3_cropped_no-bg_mask.jpg -------------------------------------------------------------------------------- /Assets/U_2_NETP_Results1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shreyas-bk/u2netdemo/3958d19a822115c6e39520e38e0303c11f299d44/Assets/U_2_NETP_Results1.jpg -------------------------------------------------------------------------------- /Assets/U_2_NETP_Results2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shreyas-bk/u2netdemo/3958d19a822115c6e39520e38e0303c11f299d44/Assets/U_2_NETP_Results2.jpg -------------------------------------------------------------------------------- /Assets/U_2_NETP_Results3.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shreyas-bk/u2netdemo/3958d19a822115c6e39520e38e0303c11f299d44/Assets/U_2_NETP_Results3.jpg -------------------------------------------------------------------------------- /Assets/u2netexamples.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shreyas-bk/u2netdemo/3958d19a822115c6e39520e38e0303c11f299d44/Assets/u2netexamples.png -------------------------------------------------------------------------------- /DEMOS/U_2_Netp_Cropper_Colab.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "nbformat": 4, 3 | "nbformat_minor": 0, 4 | "metadata": { 5 | "colab": { 6 | "name": "U_2_Netp_Cropper_Colab.ipynb", 7 | "private_outputs": true, 8 | "provenance": [], 9 | "collapsed_sections": [], 10 | "toc_visible": true 11 | }, 12 | "kernelspec": { 13 | "name": "python3", 14 | "display_name": "Python 3" 15 | }, 16 | "accelerator": "GPU" 17 | }, 18 | "cells": [ 19 | { 20 | "cell_type": "markdown", 21 | "metadata": { 22 | "id": "5t13gOcRMmsC" 23 | }, 24 | "source": [ 25 | "# U-2-NETp Cropper\n", 26 | "\n", 27 | " U-2-NET Paper: [U2-Net: Going Deeper with Nested U-Structure for Salient Object Detection](https://arxiv.org/abs/2005.09007)\n", 28 | " \n", 29 | " Repo: https://github.com/shreyas-bk/U-2-Net\n", 30 | "\n", 31 | " Original Repo: [U-2-Net Github repo](https://github.com/NathanUA/U-2-Net)\n", 32 | "\n", 33 | "References: X. Qin, Z. Zhang, C. Huang, M. Dehghan, O. R. Zaiane, and M. Jagersand, “U2-net: Going deeper with nested u-structure for salient object\n", 34 | "detection,” Pattern Recognition, vol. 106, p. 107404, 2020" 35 | ] 36 | }, 37 | { 38 | "cell_type": "code", 39 | "metadata": { 40 | "id": "yC3pYV31MezF" 41 | }, 42 | "source": [ 43 | "# make sure cloned code is saved in content\n", 44 | "%cd /content\n", 45 | "\n", 46 | "# verify CUDA\n", 47 | "!/usr/local/cuda/bin/nvcc --version\n", 48 | "\n", 49 | "# clone modified version of U-2-Net\n", 50 | "!git clone https://github.com/shreyas-bk/U-2-Net\n", 51 | "\n", 52 | "# make images directory (to store imput images) and results (to store output images) in U-2-Net folder\n", 53 | "%cd /content/U-2-Net\n", 54 | "print('making images directory')\n", 55 | "!mkdir images\n", 56 | "print('making results directory')\n", 57 | "!mkdir results\n", 58 | "print('making cropped results directory')\n", 59 | "!mkdir cropped_results\n", 60 | "\n", 61 | "#import required modules\n", 62 | "print('importing...')\n", 63 | "from google.colab import files\n", 64 | "import os\n", 65 | "from tensorflow.keras.preprocessing.image import load_img\n", 66 | "from tensorflow.keras.preprocessing.image import img_to_array\n", 67 | "import numpy as np\n", 68 | "from PIL import Image as Img\n", 69 | "import cv2\n", 70 | "print('Done!')" 71 | ], 72 | "execution_count": null, 73 | "outputs": [] 74 | }, 75 | { 76 | "cell_type": "markdown", 77 | "metadata": { 78 | "id": "i5zn3Ib3QFJa" 79 | }, 80 | "source": [ 81 | "# Make sure runtype is GPU\n", 82 | "**Runtime -> Change Runtime Type -> Hardware Accelerator -> GPU**" 83 | ] 84 | }, 85 | { 86 | "cell_type": "code", 87 | "metadata": { 88 | "id": "CEDlFFvONi5g" 89 | }, 90 | "source": [ 91 | "# change to images directory to upload image files\n", 92 | "%cd /content/U-2-Net/images\n", 93 | "uploaded = files.upload()\n" 94 | ], 95 | "execution_count": null, 96 | "outputs": [] 97 | }, 98 | { 99 | "cell_type": "code", 100 | "metadata": { 101 | "id": "-ycX9cB-jdC9" 102 | }, 103 | "source": [ 104 | "# change back to U-2-Net directory\n", 105 | "%cd /content/U-2-Net" 106 | ], 107 | "execution_count": null, 108 | "outputs": [] 109 | }, 110 | { 111 | "cell_type": "code", 112 | "metadata": { 113 | "id": "YCG6zfBhNzAH" 114 | }, 115 | "source": [ 116 | "# run the test script, and outputs are saved to results folder\n", 117 | "!python -W ignore u2net_test.py" 118 | ], 119 | "execution_count": null, 120 | "outputs": [] 121 | }, 122 | { 123 | "cell_type": "code", 124 | "metadata": { 125 | "id": "rFBODUXlk20x" 126 | }, 127 | "source": [ 128 | "import numpy as np\n", 129 | "from glob import glob\n", 130 | "# get the names of the images that were uploaded, removing .png\n", 131 | "image_dir = os.path.join(os.getcwd(), 'results/*.png')\n", 132 | "print(image_dir)\n", 133 | "file_names = glob(image_dir)\n", 134 | "#print(file_names)\n", 135 | "names = [os.path.basename(name[:-4]) for name in file_names]\n", 136 | "names" 137 | ], 138 | "execution_count": null, 139 | "outputs": [] 140 | }, 141 | { 142 | "cell_type": "code", 143 | "metadata": { 144 | "id": "hhwXTgx6e3mY" 145 | }, 146 | "source": [ 147 | "def process_image_named(name, threshold_cutoff = 0.90, use_transparency = False):\n", 148 | " result_img = load_img('/content/U-2-Net/results/'+name+'.png')\n", 149 | " # convert result-image to numpy array and rescale(255 for RBG images)\n", 150 | " RESCALE = 255\n", 151 | " out_img = img_to_array(result_img)\n", 152 | " out_img /= RESCALE\n", 153 | " # define the cutoff threshold below which, background will be removed.\n", 154 | " THRESHOLD = threshold_cutoff\n", 155 | "\n", 156 | " # refine the output\n", 157 | " out_img[out_img > THRESHOLD] = 1\n", 158 | " out_img[out_img <= THRESHOLD] = 0\n", 159 | "\n", 160 | " if use_transparency:\n", 161 | " # convert the rbg image to an rgba image and set the zero values to transparent\n", 162 | " shape = out_img.shape\n", 163 | " a_layer_init = np.ones(shape = (shape[0],shape[1],1))\n", 164 | " mul_layer = np.expand_dims(out_img[:,:,0],axis=2)\n", 165 | " a_layer = mul_layer*a_layer_init\n", 166 | " rgba_out = np.append(out_img,a_layer,axis=2)\n", 167 | " mask_img = Img.fromarray((rgba_out*RESCALE).astype('uint8'), 'RGBA')\n", 168 | " else:\n", 169 | " mask_img = Img.fromarray((out_img*RESCALE).astype('uint8'), 'RGB')\n", 170 | "\n", 171 | " # load and convert input to numpy array and rescale(255 for RBG images)\n", 172 | " input = load_img('/content/U-2-Net/images/'+name+'.jpg')\n", 173 | " inp_img = img_to_array(input)\n", 174 | " inp_img /= RESCALE\n", 175 | "\n", 176 | " \n", 177 | " if use_transparency:\n", 178 | " # since the output image is rgba, convert this also to rgba, but with no transparency\n", 179 | " a_layer = np.ones(shape = (shape[0],shape[1],1))\n", 180 | " rgba_inp = np.append(inp_img,a_layer,axis=2)\n", 181 | "\n", 182 | " #simply multiply the 2 rgba images to remove the backgound\n", 183 | " rem_back = (rgba_inp*rgba_out)\n", 184 | " rem_back_scaled = Img.fromarray((rem_back*RESCALE).astype('uint8'), 'RGBA')\n", 185 | " else:\n", 186 | " rem_back = (inp_img*out_img)\n", 187 | " rem_back_scaled = Img.fromarray((rem_back*RESCALE).astype('uint8'), 'RGB')\n", 188 | "\n", 189 | " # select a layer(can be 0,1 or 2) for bounding box creation and salient map\n", 190 | " LAYER = 2\n", 191 | " out_layer = out_img[:,:,LAYER]\n", 192 | "\n", 193 | " # find the list of points where saliency starts and ends for both axes\n", 194 | " x_starts = [np.where(out_layer[i]==1)[0][0] if len(np.where(out_layer[i]==1)[0])!=0 else out_layer.shape[0]+1 for i in range(out_layer.shape[0])]\n", 195 | " x_ends = [np.where(out_layer[i]==1)[0][-1] if len(np.where(out_layer[i]==1)[0])!=0 else 0 for i in range(out_layer.shape[0])]\n", 196 | " y_starts = [np.where(out_layer.T[i]==1)[0][0] if len(np.where(out_layer.T[i]==1)[0])!=0 else out_layer.T.shape[0]+1 for i in range(out_layer.T.shape[0])]\n", 197 | " y_ends = [np.where(out_layer.T[i]==1)[0][-1] if len(np.where(out_layer.T[i]==1)[0])!=0 else 0 for i in range(out_layer.T.shape[0])]\n", 198 | " \n", 199 | " # get the starting and ending coordinated for the box\n", 200 | " startx = min(x_starts)\n", 201 | " endx = max(x_ends)\n", 202 | " starty = min(y_starts)\n", 203 | " endy = max(y_ends)\n", 204 | " \n", 205 | " # show the resulting coordinates\n", 206 | " start = (startx,starty)\n", 207 | " end = (endx,endy)\n", 208 | " start,end\n", 209 | "\n", 210 | " cropped_rem_back_scaled = rem_back_scaled.crop((startx,starty,endx,endy))\n", 211 | " if use_transparency:\n", 212 | " cropped_rem_back_scaled.save('/content/U-2-Net/cropped_results/'+name+'_cropped_no-bg.png')\n", 213 | " else:\n", 214 | " cropped_rem_back_scaled.save('/content/U-2-Net/cropped_results/'+name+'_cropped_no-bg.jpg')\n", 215 | " \n", 216 | " cropped_mask_img = mask_img.crop((startx,starty,endx,endy))\n", 217 | "\n", 218 | " if use_transparency:\n", 219 | " cropped_mask_img.save('/content/U-2-Net/cropped_results/'+name+'_cropped_no-bg_mask.png')\n", 220 | " else:\n", 221 | " cropped_mask_img.save('/content/U-2-Net/cropped_results/'+name+'_cropped_no-bg_mask.jpg')" 222 | ], 223 | "execution_count": null, 224 | "outputs": [] 225 | }, 226 | { 227 | "cell_type": "code", 228 | "metadata": { 229 | "id": "DXgtai3Vi0J5" 230 | }, 231 | "source": [ 232 | "#Remove BG, Crop and save each image pair\n", 233 | "for name in names:\n", 234 | " process_image_named(name) #jpg, no alpha\n", 235 | " #process_image_named(name, 0.9, True) #png, with transparency" 236 | ], 237 | "execution_count": null, 238 | "outputs": [] 239 | }, 240 | { 241 | "cell_type": "markdown", 242 | "metadata": { 243 | "id": "M87UaFRq5W6x" 244 | }, 245 | "source": [ 246 | "# Zip and Download" 247 | ] 248 | }, 249 | { 250 | "cell_type": "code", 251 | "metadata": { 252 | "id": "AcA_rbqCLJJN" 253 | }, 254 | "source": [ 255 | "%cd /content/U-2-Net/cropped_results/\n", 256 | "!zip /content/cropped_results.zip ./*" 257 | ], 258 | "execution_count": null, 259 | "outputs": [] 260 | }, 261 | { 262 | "cell_type": "code", 263 | "metadata": { 264 | "id": "JSOhnNb2Obgi" 265 | }, 266 | "source": [ 267 | "# download the result\n", 268 | "# takes a while for the progress indicator to appear\n", 269 | "files.download('/content/cropped_results.zip')" 270 | ], 271 | "execution_count": null, 272 | "outputs": [] 273 | } 274 | ] 275 | } -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 shreyas-bk 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # U-2-Net-Demo 2 | Demonstration using Google Colab to show how U-2-NET can be used for Background Removal, Changing Backgrounds, Bounding Box Creation, Salient Feature Highlighting and Salient Object Cropping 3 | 4 | *Please let me know if you would like to collaborate on building a website/web demo for this repo or any other contributions* 5 | 6 | ## Check out these articles making use of this repository: 7 | [Medium: Large Scale Background Removal Using AI](https://medium.com/@shreykera7/large-scale-background-removal-using-ai-17d245b2cdb5) 8 | 9 | [I REMOVED THE BACKGROUNDS FROM 11,782 JPGS FOR FREE! USING AI!](https://drivemarketing.ca/en/blog/i-removed-the-backgrounds-from-11-782-jpgs-for-free-using-ai/) 10 | 11 | 12 | # UPDATES: upcoming update - notebook for inferencing a large number of images 13 | - [Jan 11, 2021] - Salient Object Cropping notebook is now available following this [PR](https://github.com/shreyas-bk/U-2-Net/pull/1), thanks to the contribution from [**Norod**](https://github.com/Norod)! 14 | - [Jan 05, 2021] - **Keras .h5 saved models available** in new [repository](https://github.com/shreyas-bk/U-2-Net-Keras) following this [issue](https://github.com/shreyas-bk/U-2-Net-Demo/issues/9) *Note: if there are any new requests for a TensorFlow saved model, I will upload* 15 | - [Dec 24, 2020] - change background in step by step file: [commit](https://github.com/shreyas-bk/U-2-Net-Demo/commit/71c0c8ba726bcfa1c97ddb513265ef9deede1df1) 16 | - [Oct 29, 2020] - implemented quick fix for ipynb.checkpoints error: [commit](https://github.com/shreyas-bk/U-2-Net-Demo/commit/f059e696e0ce6f6512486fbf6e9237700b710987) 17 | 18 | # Link: [Open Google Colab Notebook](https://colab.research.google.com/github/shreyas-bk/U-2-Net-Demo/blob/master/DEMOS/U_2_Netp_Demonstration_Colab.ipynb) 19 | 20 | **Click this link for step-by-step instructions:** [Open Google Colab Notebook](https://colab.research.google.com/github/shreyas-bk/U-2-Net-Demo/blob/master/DEMOS/U_2_Netp_Step_by_Step_Demonstration_Colab_v2.ipynb) {Use this to save your results too} {update: now with code to change background} 21 | 22 | **Click this link for Salient Object Cropping:** [U_2_Netp_Cropper_Colab](https://colab.research.google.com/github/shreyas-bk/U-2-Net-Demo/blob/master/DEMOS/U_2_Netp_Cropper_Colab.ipynb) 23 | 24 | # Example Results 25 | 26 | Cat: 27 | ![](Assets/U_2_NETP_Results1.jpg) 28 | 29 | Formula 1 Car: 30 | ![](Assets/U_2_NETP_Results2.jpg) 31 | 32 | Rick Astley: 33 | ![](Assets/U_2_NETP_Results3.jpg) 34 | 35 | Background Changed: 36 | 37 | 38 | 39 | Bacground Removed and Image Cropped (Mask and Resulting Image): 40 | 41 | 42 | 43 | 44 | # Useful material 45 | 46 | **U-2-NET Paper:** [U2-Net: Going Deeper with Nested U-Structure for Salient Object Detection](https://arxiv.org/abs/2005.09007) 47 | 48 | **Original Repo:** [U-2-Net Github repo](https://github.com/NathanUA/U-2-Net) 49 | 50 | **Modified repo for that this colab uses:** [Modified fork](https://github.com/shreyas-bk/U-2-Net) 51 | 52 | **References:** X. Qin, Z. Zhang, C. Huang, M. Dehghan, O. R. Zaiane, and M. Jagersand, “U2-net: Going deeper with nested u-structure for salient object 53 | detection,” Pattern Recognition, vol. 106, p. 107404, 2020 54 | 55 | # Other Results 56 | 57 | ![](Assets/u2netexamples.png) 58 | 59 | 60 | # Details about U-2-Net 61 | The following is an excerpt from the paper: 62 | *'In this paper, we design a simple yet powerful deep network architecture, U2-Net, for salient object detection (SOD). The architecture of our U2-Net is a two-level nested U-structure. The design has the following advantages: (1) it is able to capture more contextual information from different scales thanks to the mixture of receptive fields of different sizes in our proposed ReSidual U-blocks (RSU), (2) it increases the depth of the whole architecture without significantly increasing the computational cost because of the pooling operations used in these RSU blocks. This architecture enables us to train a deep network from scratch without using backbones from image classification tasks. We instantiate two models of the proposed architecture, U2-Net (176.3 MB, 30 FPS on GTX 1080Ti GPU) and U2-Net† (4.7 MB, 40 FPS), to facilitate the usage in different environments. Both models achieve competitive performance on six SOD datasets.'* 63 | 64 | # U2NET Model 65 | 66 | ![](Assets/U2NETPRmodel.png) 67 | 68 | TODO: 69 | - remove dependancy to secondary repo 70 | - algorithmic post processing 71 | - flask demo 72 | - multiproc cpu code 73 | - benchmarks 74 | - video background removal 75 | - change backgrounds for main notebook 76 | - support .jpeg images 77 | - upload python code for webcam support 78 | - add demo for portrait feature of U-2-Net 79 | - add license to nbs 80 | - maybe make a webpage for bounding box creation over large datasets 81 | 82 | **Please feel free to suggest changes to improve this repository** 83 | --------------------------------------------------------------------------------