26 lines
861 B
Python
26 lines
861 B
Python
import tensorflow as tf
|
|
import numpy as np
|
|
# from keras import utils
|
|
# from tensorflow.keras.utils.generic_utils import get_custom_objects
|
|
from yololoss import YOLOLoss
|
|
|
|
class YOLOActivation(tf.keras.layers.Layer):
|
|
def __init__(self, name = 'yolo_activation', **kwargs):
|
|
super(YOLOActivation, self).__init__(name = name)
|
|
|
|
def call(self, inputs):
|
|
conf = tf.sigmoid(inputs[..., 0])
|
|
conf = tf.expand_dims(conf, axis = -1)
|
|
|
|
coord = tf.sigmoid(inputs[..., 1:5])
|
|
|
|
class_probs = tf.nn.softmax(inputs[..., 5:])
|
|
|
|
outputs = tf.concat([conf, coord, class_probs], axis = -1)
|
|
|
|
return outputs
|
|
|
|
tf.keras.utils.get_custom_objects().update({'YOLOActivation': YOLOActivation})
|
|
tf.keras.utils.get_custom_objects().update({'YOLOLoss': YOLOLoss})
|
|
|
|
model = tf.keras.models.load_model('model/SIBIDetection_YOLOv2_780_2.h5') |