segmentation metric with tensorflow

segm_metric_testcase.py

"""
Created on Wed Oct 25 13:44:51 2017

@author: dlituiev
"""
import numpy as np
import keras
from keras import backend as K
from keras.layers import (InputLayer, Conv2D, Dense, Activation,
                          AveragePooling2D, GlobalAveragePooling2D)
from keras.applications.inception_v3 import InceptionV3
import tensorflow as tf
from PIL import Image
from keras.preprocessing.image import ImageDataGenerator

def metric_per_channel_tf(label, prediction, nch=3, metric=tf.metrics.accuracy):
    prec = []
    #label = tf.stack([label])
    print("label.shape", label.shape)
    print("prediction.shape", prediction.shape)
    dummy = [0] * len(prediction.shape[:-1])
    shape = ([(x.value if x.value is not None else -1) for x in prediction.shape[:-1]] + [1])
    print("SHAPE:", shape)
    for cc in range(nch):
        print("start", (dummy + [cc]))
        print("end", shape )

        pred_channel = tf.slice(prediction, (dummy + [cc]), shape)
        pred_channel = tf.reshape(pred_channel, shape[:-1])
        # label_channel = tf.equal(label,cc)
        label_channel = label[:,:,:,cc]
        _, prec_ = metric(label_channel, pred_channel)
        prec.append(prec_)
    prec = tf.stack(prec)
    return tf.reduce_sum(prec)
    
# Model

model = keras.models.Sequential()
model.add(InputLayer((299,299,3)))
model.add(Conv2D(8,(3,3)))
model.add(AveragePooling2D(3,3))
model.add(Conv2D(4,(3,3)))
model.add(AveragePooling2D(3,3))
model.add(Conv2D(2,(3,3)))
model.add(Activation("softmax"))

print(model.output.shape)

# Data
x = np.zeros((4,299,299,3))
for ii in range(len(x)):
    x_ = x[ii,:,:,0].copy()
    start = np.random.randint(100, 200, size=(2,1))
    path = (start + np.cumsum(np.random.randint(-1,2, size=(2,100)), axis=1)) % 299
    path_ = np.ravel_multi_index(path, dims=(299,299))
    x_.ravel()[path_] = 256
    for cc in range(x.shape[-1]):
        x[ii,:,:,cc] = x_

y = np.stack(
        [x[:,::10,::10,0] > 0,
         x[:,::10,::10,0] == 0,
         ], axis=-1)
         

config = tf.ConfigProto(log_device_placement=True)
sess0 = tf.Session(config=config)

init_g = tf.global_variables_initializer()
init_l = tf.local_variables_initializer()

with sess0.as_default() as sess:
    def accuracy_per_channel(x,y):
        return metric_per_channel_tf(x,y, nch=2, metric=tf.metrics.accuracy)

    sess.run(init_g)
    sess.run(init_l)

    model.compile(optimizer='Adam',
                  loss='categorical_crossentropy',
                  metrics=[accuracy_per_channel],
                  )
    model.fit(x, y)
    print("DONE")