text: # -*- coding: utf-8 -*-
  """“lhj semantic segmentation.ipynb”
  
  Automatically generated by Colaboratory.
  
  Original file is located at
      https://colab.research.google.com/drive/18MQ_vmrqz_O7-9hx-KhwM6C9MIGY3SFA
  
  ## Import Libraries
  """
  
  # Commented out IPython magic to ensure Python compatibility.
  import os
  from io import BytesIO
  import tarfile
  import tempfile
  from six.moves import urllib
  
  from matplotlib import gridspec
  from matplotlib import pyplot as plt
  import numpy as np
  from PIL import Image
  from scipy import ndimage
  # %tensorflow_version 1.x
  import tensorflow as tf
  
  
  class DeepLabModel(object):
      """Class to load deeplab model and run inference."""
  
      INPUT_TENSOR_NAME = 'ImageTensor:0'
      OUTPUT_TENSOR_NAME = 'SemanticPredictions:0'
      INPUT_SIZE = 513
      FROZEN_GRAPH_NAME = 'frozen_inference_graph'
  
      def __init__(self, tarball_path):
          """Creates and loads pretrained deeplab model."""
          self.graph = tf.Graph()
  
          graph_def = None
          # Extract frozen graph from tar archive.
          tar_file = tarfile.open(tarball_path)
          for tar_info in tar_file.getmembers():
              if self.FROZEN_GRAPH_NAME in os.path.basename(tar_info.name):
                  file_handle = tar_file.extractfile(tar_info)
                  graph_def = tf.GraphDef.FromString(file_handle.read())
                  break
  
          tar_file.close()
  
          if graph_def is None:
              raise RuntimeError('Cannot find inference graph in tar archive.')
  
          with self.graph.as_default():
              tf.import_graph_def(graph_def, name='')
  
          self.sess = tf.Session(graph=self.graph)
  
      def run(self, image):
          """Runs inference on a single image.
  
          Args:
            image: A PIL.Image object, raw input image.
  
          Returns:
            resized_image: RGB image resized from original input image.
            seg_map: Segmentation map of `resized_image`.
          """
          width, height = image.size
          original_size = [width, height]
          resize_ratio = 1.0 * self.INPUT_SIZE / max(width, height)
          target_size = (int(resize_ratio * width), int(resize_ratio * height))
          resized_image = image.convert('RGB').resize(target_size, Image.ANTIALIAS)
          batch_seg_map = self.sess.run(
              self.OUTPUT_TENSOR_NAME,
              feed_dict={self.INPUT_TENSOR_NAME: [np.asarray(resized_image)]})
          seg_map = batch_seg_map[0]
          return resized_image, seg_map, original_size
  
  
  '''
    |= : in-place or operation. find the or between two sets and reolace the first set
  '''
  
  
  def create_pascal_label_colormap():
      """Creates a label colormap used in PASCAL VOC segmentation benchmark.
  
        主要是得到 分割后每个对象的颜色映射
      Returns:
        A Colormap for visualizing segmentation results.
      """
      colormap = np.zeros((256, 3), dtype=int)
      ind = np.arange(256, dtype=int)
  
      for shift in reversed(range(8)):
          for channel in range(3):
              colormap[:, channel] |= ((ind >> channel) & 1) << shift
          ind >>= 3
  
      return colormap
  
  
  def label_to_color_image(label):
      """Adds color defined by the dataset colormap to the label.
  
      Args:
        label: A 2D array with integer type, storing the segmentation label.
  
      Returns:
        result: A 2D array with floating type. The element of the array
          is the color indexed by the corresponding element in the input label
          to the PASCAL color map.
  
      Raises:
        ValueError: If label is not of rank 2 or its value is larger than color
          map maximum entry.
      """
      if label.ndim != 2:
          raise ValueError('Expect 2-D input label')
  
      colormap = create_pascal_label_colormap()
  
      if np.max(label) >= len(colormap):
          raise ValueError('label value too large.')
  
      return colormap[label]
  
  
  def find_people(seg_map):
      """
      get the pixel of people
      the output of the network is 0—20 label in each pixel.  (represent an object)
      np.putmask  function: to use mask to find people
      """
      x = seg_map
      np.putmask(x, x != 15, 0)  # 15 represents person
      y = np.expand_dims(x, 2).repeat(3,
                                      axis=2)  # expand the dimensions, seg_map (axis_x, axis_y, 1)  --->  (axis_x, axis_y, 3)
      # x = x/(15*255)
      y = y / 15
      return y, x
  
  
  def vis_segmentation(image, seg_map, original_size):
      """Visualizes input image, segmentation map and overlay view."""
      plt.figure(figsize=(15, 5))
      grid_spec = gridspec.GridSpec(1, 5, width_ratios=[6, 6, 6, 6, 1])
  
      plt.subplot(grid_spec[0])
      plt.imshow(image)
      plt.axis('off')
      plt.title('input image')
      # save the input image
      PIL_img = Image.fromarray(np.uint8(image)).convert('RGB')
      PIL_img = PIL_img.resize((original_size[0], original_size[1]))
      PIL_img.save('input.png')
  
      plt.subplot(grid_spec[1])
      seg_image = label_to_color_image(seg_map).astype(np.uint8)
      plt.imshow(seg_image)
      plt.axis('off')
      plt.title('segmentation map')
  
      plt.subplot(grid_spec[2])
      plt.imshow(image)
      plt.imshow(seg_image, alpha=0.7)
      plt.axis('off')
      plt.title('segmentation overlay')
  
      plt.subplot(grid_spec[3])
      seg_people, alpha_seg = find_people(seg_map)
      seg_image = np.multiply(seg_people, image)
      seg_image_pil = np.multiply(seg_people / 255, image)
      plt.imshow(seg_image_pil)
      plt.axis('off')
      plt.title('only get people')
      # save the output segmentation image
      PIL_seg = Image.fromarray(np.uint8(seg_image)).convert('RGB')
      tri_map = alpha_seg / 15 * 255
      PIL_seg_alpha = Image.fromarray(np.uint8(tri_map))
      width, height = PIL_seg.size
      ori_PIL_seg = PIL_seg.resize((original_size[0], original_size[1]))
      ori_PIL_seg_alpha = PIL_seg_alpha.resize((original_size[0], original_size[1]))
      ori_PIL_seg.save('seg.png')
      ori_PIL_seg_alpha.save('seg_a.png')
  
      unique_labels = np.unique(seg_map)
      ax = plt.subplot(grid_spec[4])
      plt.imshow(
          FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation='nearest')
      ax.yaxis.tick_right()
      plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])
      plt.xticks([], [])
      ax.tick_params(width=0.0)
      plt.grid('off')
  
      plt.show()
  
  def output_seg_map(image, seg_map, original_size):
      seg_people, alpha_seg = find_people(seg_map)
      seg_image = np.multiply(seg_people, image)
      PIL_seg = Image.fromarray(np.uint8(seg_image)).convert('RGB')
      alpha_seg = alpha_seg / 15 * 255
      PIL_seg_alpha = Image.fromarray(np.uint8(alpha_seg))
  
      ori_PIL_seg = PIL_seg.resize((original_size[0], original_size[1]))
      ori_PIL_seg_alpha = PIL_seg_alpha.resize((original_size[0], original_size[1]))
  
      alpha = ori_PIL_seg_alpha
      trimap = generate_trimap(alpha)
      trimap = Image.fromarray(trimap)
      return ori_PIL_seg, ori_PIL_seg_alpha, trimap
  
  
  
  
  
  
  LABEL_NAMES = np.asarray([
      'background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus',
      'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike',
      'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tv'
  ])
  
  FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
  FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)
  
  """## Select a pretrained model
  We have trained the DeepLab model using various backbone networks. Select one from the MODEL_NAME list.
  """
  
  MODEL_NAME = 'mobilenetv2_coco_voctrainval'  # @param ['mobilenetv2_coco_voctrainaug', 'mobilenetv2_coco_voctrainval', 'xception_coco_voctrainaug', 'xception_coco_voctrainval']
  
  _DOWNLOAD_URL_PREFIX = 'http://download.tensorflow.org/models/'
  _MODEL_URLS = {
      'mobilenetv2_coco_voctrainaug':
          'deeplabv3_mnv2_pascal_train_aug_2018_01_29.tar.gz',
      'mobilenetv2_coco_voctrainval':
          'deeplabv3_mnv2_pascal_trainval_2018_01_29.tar.gz',
      'xception_coco_voctrainaug':
          'deeplabv3_pascal_train_aug_2018_01_04.tar.gz',
      'xception_coco_voctrainval':
          'deeplabv3_pascal_trainval_2018_01_04.tar.gz',
  }
  _TARBALL_NAME = 'deeplab_model.tar.gz'
  
  '''
  ## already downloaded the pretrained model
  
  ## model_dir = tempfile.mkdtemp()
  model_dir = 'E:/MSRA/code/semantic segmentation/model'
  tf.gfile.MakeDirs(model_dir)
  
  download_path = os.path.join(model_dir, _TARBALL_NAME)
  print('downloading model, this might take a while...')
  urllib.request.urlretrieve(_DOWNLOAD_URL_PREFIX + _MODEL_URLS[MODEL_NAME],
                     download_path)
  print('download completed! loading DeepLab model...')
  
  MODEL = DeepLabModel(download_path)
  print('model loaded successfully!')
  '''
  
  model_dir = 'E:/MSRA/code/semantic segmentation/model'
  tf.gfile.MakeDirs(model_dir)
  model_path = os.path.join(model_dir, _TARBALL_NAME)
  MODEL = DeepLabModel(model_path)
  
  """## Run on sample images
  
  Select one of sample images (leave `IMAGE_URL` empty) or feed any internet image
  url for inference.
  
  Note that this colab uses single scale inference for fast computation,
  so the results may slightly differ from the visualizations in the
  [README](https://github.com/tensorflow/models/blob/master/research/deeplab/README.md) file,
  which uses multi-scale and left-right flipped inputs.
  """
  
  
  def generate_trimap(alpha):
      fg = np.array(np.equal(alpha, 255).astype(np.float32))
      unknown = np.array(np.not_equal(alpha, 0).astype(np.float32))  # unknown = alpha > 0
      unknown = unknown - fg
      unknown = ndimage.morphology.distance_transform_edt(unknown == 0) <= np.random.randint(1, 13)
      trimap = fg * 255
      trimap[unknown] = 128
      return trimap.astype(np.uint8)
  
  
  def run_segmentation_trimap(image_base, pic_name, directory):
      """ implement semantic segmentation on local picture """
      for i in pic_name:
          name = i
          im_path = image_base + directory[0] + name + '.png'
          try:
              original_im = Image.open(im_path)
          except IOError:
              print('cannot open image')
              return
          print('running deeplab on image %s...' % im_path)
          resized_im, seg_map, original_size = MODEL.run(original_im)
          ori_PIL_seg, ori_PIL_seg_alpha, trimap = output_seg_map(resized_im, seg_map, original_size)
  
          ori_seg_path = image_base + directory[1] + name + '.png'
          ori_alpha_path = image_base + directory[2] + name + '.png'
          trimap_path = image_base + directory[3] + name + '.png'
          ori_PIL_seg.save(ori_seg_path)
          ori_PIL_seg_alpha.save(ori_alpha_path)
          trimap.save(trimap_path)
  
  
  
  """
  running model and visualize segmentation pictures
  image_url = IMAGE_URL or _SAMPLE_URL % SAMPLE_IMAGE
  run_visualization(image_url)
  """
  
  
  
  
  if __name__ == "__main__":
      pic_name = ['man1', 'man2', 'man3', 'man4', 'man5', 'man6', 'woman1', 'woman2']
      image_base = 'E:/MSRA/code/data/'
      directory = ['people/', 'people_seg/','people_alpha/','people_trimap/']
      run_segmentation_trimap(image_base, pic_name, directory)
  
  
  