基于tensorflow.keras的centernet复现

 

from functools import partial
from tensorflow.keras.layers import Conv2D, Add, ZeroPadding2D, UpSampling2D, concatenate, MaxPooling2D, Lambda, add
from tensorflow.keras.layers import LeakyReLU,ReLU
from tensorflow.keras.layers import BatchNormalization
from tensorflow.keras.regularizers import l2
import sys
sys.path.append("./")
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Input
from tensorflow.keras.models import Model
from tensorflow.keras.applications import MobileNetV2
import tensorflow.keras.backend as K
from resnet50 import ResNet50

def losses(y_true,y_pred):
    hm = y_true[...,:-4]
    wh = y_true[...,-4:-2]
    reg = y_true[...,-2:]
    y_true = {'hm':hm,'wh':wh,'reg':reg}
    hm_output =  y_pred[...,:-4]
    wh_output =  y_pred[...,-4:-2]
    reg_output = y_pred[...,-2:]
    output ={ 'hm':hm_output,'wh':wh_output,'reg':reg_output}
    mask = tf.equal(y_true['hm'],1)
    mask = tf.cast(mask,tf.float32)
    mask_reg = tf.reduce_max(mask,axis = -1)
    mask_reg = tf.expand_dims(mask_reg,-1)
    mask_reg = tf.concat([mask_reg,mask_reg],-1)
    N = tf.reduce_sum(mask,1)
    N = tf.reduce_sum(N,1)
    N = tf.reduce_sum(N,1)
    #output['hm'] = tf.nn.sigmoid(output['hm'])
    loss_hm_pos = -1.0*tf.pow(1.-output['hm'],2.)*tf.log(output['hm']+1e-12) * mask
    loss_hm_neg = -1.0*tf.pow(1.-y_true['hm'],4)*tf.pow(output['hm'],2)*tf.log(1.-output['hm']+1e-12)*(1.-mask)
    loss_hm = tf.reduce_sum(loss_hm_pos+loss_hm_neg,axis=1)
    loss_hm = tf.reduce_sum(loss_hm,axis=1)
    loss_hm = tf.reduce_sum(loss_hm,axis=1)/N
    loss_wh = tf.abs(y_true['wh']-output['wh']) * mask_reg
    loss_wh = tf.reduce_sum(loss_wh,axis=1)
    loss_wh = tf.reduce_sum(loss_wh,axis=1)
    loss_wh = tf.reduce_sum(loss_wh,axis=1)/N
    loss_reg = tf.abs(y_true['reg']-output['reg']) * mask_reg
    loss_reg = tf.reduce_sum(loss_reg,axis=1)
    loss_reg = tf.reduce_sum(loss_reg,axis=1)
    loss_reg = tf.reduce_sum(loss_reg,axis=1)/N
    loss_total =loss_hm+1.0*loss_wh+loss_reg
    return loss_hm,loss_wh,loss_reg,loss_total



#--------------------------------------------------#
#   GET_C获取featurmap
#--------------------------------------------------#
def GET_C(input_shape=[256,256,3],classes=80):
    return ResNet50(input_shape,classes)

#--------------------------------------------------#
#   FPN
#--------------------------------------------------#
def FPN(C2,C3,C4,C5):
    P5 = Conv2D(256,1,padding='same')(C5)
    P4 = add([UpSampling2D()(P5),Conv2D(256,1,padding='same')(C4)])
    P3 = add([UpSampling2D()(P4),Conv2D(256,1,padding='same')(C3)])
    P2 = add([UpSampling2D()(P3),Conv2D(256,1,padding='same')(C2)])
    P6 = Conv2D(256,3,padding='same')(P5)
    return  [P2,P3,P4,P5,P6]


def merge_layers(big_layer,small_layer):
    size = int(big_layer.shape[1])//int(small_layer.shape[1])
    n_filter=int(big_layer.shape[-1])
    big_layer_conv =  Conv2D(n_filter,3,padding='same')(big_layer)  
    small_layer_up = UpSampling2D(size=(size,size))(small_layer)
    small_layer_up = Conv2D(n_filter,3,padding='same')(small_layer_up)  
    out = concatenate([big_layer_conv,small_layer_up],axis=-1)
    out =  Conv2D(n_filter,3,padding='same')(out)  
    out = BatchNormalization()(out)
    out = LeakyReLU(0.2)(out)
    return out
    

def center_branch(feat,n_classes=9):
    #hm
    hm = Conv2D(256,3,padding='same')(feat)   #slim.conv2d(feat,256,[3,3])
    hm = ReLU()(hm)
    hm = Conv2D(n_classes,1,padding='same',activation='sigmoid')(hm)  #slim.conv2d(hm,n_classes,[1,1],activation_fn=tf.sigmoid)
    #WH
    wh = Conv2D(256,3,padding='same')(feat)
    wh = ReLU()(wh)
    wh = Conv2D(2,1,padding='same')(hm)  #slim.conv2d(wh,2,[1,1],activation_fn=None)
    #reg
    reg = Conv2D(256,3,padding='same')(feat)
    reg = ReLU()(reg)
    reg = Conv2D(2,1,padding='same')(hm)  #slim.conv2d(reg,2,[1,1],activation_fn=None)
    output = concatenate([hm,wh,reg],axis=-1)
    return output#[hm,wh,reg]#{'hm':hm,'wh':wh,'reg':reg}


#--------------------------------------------------#
#   build_model
#--------------------------------------------------#
def build_model(input_shape=[256,256,3],classes=10):
    img_input = Input(shape=input_shape)
    
    # 获取res50特征层
    C2,C3,C4,C5 = GET_C(img_input,classes)
    #DLA-34过程
    #stage1
    C2_1 = merge_layers(C2,C3)
    C3_1 = merge_layers(C3,C4)
    C4_1 = merge_layers(C4,C5)

    #stage2
    C2_2 = merge_layers(C2_1,C3_1)
    C3_2 = merge_layers(C3_1,C4_1)
    
    #stage3
    C2_3 =merge_layers(C2_2,C3_2)
    C3_3 = merge_layers(C2_3,C3_2)
    C4_3 = merge_layers(C3_3,C4_1)
    C5_3 = merge_layers(C4_3,C5)
    C5_3 = UpSampling2D()(C5_3)
    result = center_branch(C5_3,n_classes=10)
    label_input= Input(shape=result.shape[1:])

    # 添加损失函数layer
    loss_hm,loss_wh,loss_reg,loss_total =Lambda(lambda x:losses(*x))([label_input,result])
    # 给损失layer命名
    loss_hm = Lambda(lambda x:x,name='hm_loss')(loss_hm)
    loss_wh = Lambda(lambda x:x,name='wh_loss')(loss_wh)
    loss_reg = Lambda(lambda x:x,name='reg_loss')(loss_reg)
    loss_total = Lambda(lambda x:x,name='total_loss')(loss_total)

    model = Model(inputs= [img_input,label_input],outputs = [result, loss_hm,loss_wh,loss_reg,loss_total])                                                           
    model.add_loss(model.get_layer('hm_loss').output)
    model.add_loss(model.get_layer('wh_loss').output)
    model.add_loss(model.get_layer('reg_loss').output)
    model.add_loss(model.get_layer('total_loss').output)
    return model