2017-03-22 21 views
5

我正在Tensorflow中實現語義分割網絡,並試圖弄清楚如何在培訓期間寫出標籤的摘要圖像。我想以與Pascal VOC數據集中使用的class segmentation annotations類似的樣式對圖像進行編碼。Tensorflow:如何創建Pascal VOC風格圖像

例如,假設我有一個網絡,其批量大小爲1,帶有4個類。該網絡最終的預測具有塑造[1, 3, 3, 4]

基本上我想借此輸出預測,並通過argmin運行它以獲得含有最可能的類在輸出各點的張量:

[[[0, 1, 3], 
    [2, 0, 1], 
    [3, 1, 2]]] 

的註釋的圖像使用255種顏色的調色板對標籤進行編碼。我有包含所有顏色三元組的張量:

[[ 0, 0, 0], 
    [128, 0, 0], 
    [ 0, 128, 0], 
    [128, 128, 0], 
    [ 0, 0, 128], 
    ... 
    [224, 224, 192]] 

我怎樣才能獲得形狀[1, 3, 3, 3](單3×3顏色圖像)的張量的索引進使用從argmin獲得的值調色板?

[[palette[0], palette[1], palette[3]], 
[palette[2], palette[0], palette[1]], 
[palette[3], palette[1], palette[2]]] 

我可以很容易地包裹在tf.py_func一些numpy的和PIL代碼,但我不知道是否有獲得這種結果的純Tensorflow方式。

編輯: 對於那些好奇,這是我剛剛使用numpy的解決方案。它的工作原理相當不錯,但我還是不喜歡使用tf.py_func

import numpy as np 
import tensorflow as tf 


def voc_colormap(N=256): 
    bitget = lambda val, idx: ((val & (1 << idx)) != 0) 

    cmap = np.zeros((N, 3), dtype=np.uint8) 
    for i in range(N): 
     r = g = b = 0 
     c = i 
     for j in range(8): 
      r |= (bitget(c, 0) << 7 - j) 
      g |= (bitget(c, 1) << 7 - j) 
      b |= (bitget(c, 2) << 7 - j) 
      c >>= 3 

     cmap[i, :] = [r, g, b] 
    return cmap 


VOC_COLORMAP = voc_colormap() 


def grayscale_to_voc(input, name="grayscale_to_voc"): 
    return tf.py_func(grayscale_to_voc_impl, [input], tf.uint8, stateful=False, name=name) 


def grayscale_to_voc_impl(input): 
    return np.squeeze(VOC_COLORMAP[input]) 

回答

1

您可以使用tf.gather_nd(),但你將需要修改調色板和logits的形狀以獲得所需的圖像,例如:

import tensorflow as tf 
import numpy as np 
import PIL.Image as Image 

# We can load the palette from some random image in the PASCAL VOC dataset 
palette = Image.open('.../VOC2012/SegmentationClass/2007_000032.png').getpalette() 

# We build a random logits tensor of the requested size 
batch_size = 1 
height = width = 3 
num_classes = 4 
np.random.seed(1234) 
logits = np.random.random_sample((batch_size, height, width, num_classes)) 
logits_argmax = np.argmax(logits, axis=3) # shape = (1, 3, 3) 
# array([[[3, 3, 0], 
#   [1, 3, 1], 
#   [0, 2, 0]]]) 

sess = tf.InteractiveSession() 
image = tf.gather_nd(
    params=tf.reshape(palette, [-1, 3]), # reshaped from list to RGB 
    indices=tf.reshape(logits_argmax, [batch_size, -1, 1])) 
image = tf.cast(tf.reshape(image, [batch_size, height, width, 3]), tf.uint8) 
sess.run(image) 
# array([[[[128, 128, 0], 
#   [128, 128, 0], 
#   [ 0, 0, 0]], 
#   [[128, 0, 0], 
#   [128, 128, 0], 
#   [128, 0, 0]], 
#   [[ 0, 0, 0], 
#   [ 0, 128, 0], 
#   [ 0, 0, 0]]]], dtype=uint8) 

由此產生的張量可以直接送到tf.summary.image(),但取決於你的實現,你應該在摘要前對它進行上採樣。