rishizek / tensorflow-deeplab-v3-plus

DeepLabv3+ built in TensorFlow
MIT License
834 stars 307 forks source link

How can i get the mask of a particular class in inference.py? #11

Closed bruceyang2012 closed 6 years ago

bruceyang2012 commented 6 years ago

Hi, can you guide me to get de mask of person in inference.py? I need it in a hurry, great thanks to you if you can help me with it.

rishizek commented 6 years ago

Hi @bruceyang2012 , thanks for your interest in the repo. You mean you want to mask only person class and convert other classes to background?

One simple way to achieve that is following:

@@ -172,6 +172,10 @@ def deeplabv3_plus_model_fn(features, labels, mode, params):
   logits = network(features, mode == tf.estimator.ModeKeys.TRAIN)

   pred_classes = tf.expand_dims(tf.argmax(logits, axis=3, output_type=tf.int32), axis=3)
+  if mode == tf.estimator.ModeKeys.PREDICT:
+    where_parson = tf.equal(pred_classes, 15)  # Convert person class (=15) to True and other to False
+    pred_classes = tf.cast(where_parson, tf.int32)  # Convert True to 1 and False to zero
+    pred_classes = pred_classes * 15  # Convert 1 to 15

   pred_decoded_labels = tf.py_func(preprocessing.decode_labels,
                                    [pred_classes, params['batch_size'], params['num_classes']],

And the result is shown below: 2007_000129_mask which was originally: 2007_000129_mask I hope this solves your problem.

bruceyang2012 commented 6 years ago

Very thanks for such detailed answer!I become your fans now.(^-^)