├── README.md ├── sample.jpg ├── sample_bgremoved.png ├── seg.py └── setup.sh /README.md: -------------------------------------------------------------------------------- 1 | # Portrait Segmentation using Tensorflow 2 | 3 | This script removes the background from an input image. You can read more about segmentation [here](http://colab.research.google.com/github/tensorflow/models/blob/master/research/deeplab/deeplab_demo.ipynb) 4 | 5 | ### Setup 6 | The script setup.sh downloads the trained model and sets it up so that the seg.py script can understand. 7 | > ./setup.sh 8 | 9 | ### Running the script 10 | Go ahead and use the script as specified below, to execute fast but lower accuracy model: 11 | > python3 seg.py sample.jpg sample.png 12 | 13 | For better accuracy, albiet a slower approach, go ahead and try : 14 | > python3 seg.py sample.jpg sample.png 1 15 | 16 | ### Dependencies 17 | > tensorflow, PIL 18 | 19 | ### Sample Result 20 | Input: 21 | ![alt text](https://github.com/callmesusheel/image-background-removal/raw/master/sample.jpg "Input") 22 | 23 | Output: 24 | ![alt text](https://github.com/callmesusheel/image-background-removal/raw/master/sample_bgremoved.png "Output") -------------------------------------------------------------------------------- /sample.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/susheelsk/image-background-removal/6ab3c0eb248f4712102337888a29fc724dd551fc/sample.jpg -------------------------------------------------------------------------------- /sample_bgremoved.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/susheelsk/image-background-removal/6ab3c0eb248f4712102337888a29fc724dd551fc/sample_bgremoved.png -------------------------------------------------------------------------------- /seg.py: -------------------------------------------------------------------------------- 1 | import os 2 | from io import BytesIO 3 | 4 | import numpy as np 5 | from PIL import Image 6 | 7 | import tensorflow as tf 8 | import sys 9 | import datetime 10 | 11 | 12 | class DeepLabModel(object): 13 | """Class to load deeplab model and run inference.""" 14 | 15 | INPUT_TENSOR_NAME = 'ImageTensor:0' 16 | OUTPUT_TENSOR_NAME = 'SemanticPredictions:0' 17 | INPUT_SIZE = 513 18 | FROZEN_GRAPH_NAME = 'frozen_inference_graph' 19 | 20 | def __init__(self, tarball_path): 21 | """Creates and loads pretrained deeplab model.""" 22 | self.graph = tf.Graph() 23 | 24 | graph_def = None 25 | graph_def = tf.GraphDef.FromString(open(tarball_path + "/frozen_inference_graph.pb", "rb").read()) 26 | 27 | if graph_def is None: 28 | raise RuntimeError('Cannot find inference graph in tar archive.') 29 | 30 | with self.graph.as_default(): 31 | tf.import_graph_def(graph_def, name='') 32 | 33 | self.sess = tf.Session(graph=self.graph) 34 | 35 | def run(self, image): 36 | """Runs inference on a single image. 37 | 38 | Args: 39 | image: A PIL.Image object, raw input image. 40 | 41 | Returns: 42 | resized_image: RGB image resized from original input image. 43 | seg_map: Segmentation map of `resized_image`. 44 | """ 45 | start = datetime.datetime.now() 46 | 47 | width, height = image.size 48 | resize_ratio = 1.0 * self.INPUT_SIZE / max(width, height) 49 | target_size = (int(resize_ratio * width), int(resize_ratio * height)) 50 | resized_image = image.convert('RGB').resize(target_size, Image.ANTIALIAS) 51 | batch_seg_map = self.sess.run( 52 | self.OUTPUT_TENSOR_NAME, 53 | feed_dict={self.INPUT_TENSOR_NAME: [np.asarray(resized_image)]}) 54 | seg_map = batch_seg_map[0] 55 | 56 | end = datetime.datetime.now() 57 | 58 | diff = end - start 59 | print("Time taken to evaluate segmentation is : " + str(diff)) 60 | 61 | return resized_image, seg_map 62 | 63 | def drawSegment(baseImg, matImg): 64 | width, height = baseImg.size 65 | dummyImg = np.zeros([height, width, 4], dtype=np.uint8) 66 | for x in range(width): 67 | for y in range(height): 68 | color = matImg[y,x] 69 | (r,g,b) = baseImg.getpixel((x,y)) 70 | if color == 0: 71 | dummyImg[y,x,3] = 0 72 | else : 73 | dummyImg[y,x] = [r,g,b,255] 74 | img = Image.fromarray(dummyImg) 75 | img.save(outputFilePath) 76 | 77 | 78 | inputFilePath = sys.argv[1] 79 | outputFilePath = sys.argv[2] 80 | 81 | if inputFilePath is None or outputFilePath is None: 82 | print("Bad parameters. Please specify input file path and output file path") 83 | exit() 84 | 85 | modelType = "mobile_net_model" 86 | if len(sys.argv) > 3 and sys.argv[3] == "1": 87 | modelType = "xception_model" 88 | 89 | MODEL = DeepLabModel(modelType) 90 | print('model loaded successfully : ' + modelType) 91 | 92 | def run_visualization(filepath): 93 | """Inferences DeepLab model and visualizes result.""" 94 | try: 95 | print("Trying to open : " + sys.argv[1]) 96 | # f = open(sys.argv[1]) 97 | jpeg_str = open(filepath, "rb").read() 98 | orignal_im = Image.open(BytesIO(jpeg_str)) 99 | except IOError: 100 | print('Cannot retrieve image. Please check file: ' + filepath) 101 | return 102 | 103 | print('running deeplab on image %s...' % filepath) 104 | resized_im, seg_map = MODEL.run(orignal_im) 105 | 106 | # vis_segmentation(resized_im, seg_map) 107 | drawSegment(resized_im, seg_map) 108 | 109 | run_visualization(inputFilePath) 110 | 111 | -------------------------------------------------------------------------------- /setup.sh: -------------------------------------------------------------------------------- 1 | wget http://download.tensorflow.org/models/deeplabv3_mnv2_pascal_train_aug_2018_01_29.tar.gz 2 | wget http://download.tensorflow.org/models/deeplabv3_pascal_train_aug_2018_01_04.tar.gz 3 | 4 | mkdir mobile_net_model 5 | mkdir xception_model 6 | tar xvzf deeplabv3_mnv2_pascal_train_aug_2018_01_29.tar.gz -C mobile_net_model --strip=1 7 | tar xvzf deeplabv3_pascal_train_aug_2018_01_04.tar.gz -C xception_model --strip=1 8 | 9 | rm deeplabv3_mnv2_pascal_train_aug_2018_01_29.tar.gz 10 | rm deeplabv3_pascal_train_aug_2018_01_04.tar.gz --------------------------------------------------------------------------------