-
Notifications
You must be signed in to change notification settings - Fork 0
/
vgg.py
35 lines (27 loc) · 1.23 KB
/
vgg.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import tensorflow as tf
from tensorflow.keras.applications.vgg16 import VGG16, preprocess_input
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Activation, Flatten, Dense, Input, BatchNormalization
from tensorflow.keras.models import Model
class TransferModel(Model):
def __init__(self, input_shape,classes = 1, chanDim=-1):
super(TransferModel, self).__init__()
self.classes = classes
self.base_model = VGG16(input_shape = input_shape, include_top=False, weights = 'imagenet')
self.flatten = Flatten()
self.dense = Dense(256, activation = 'relu')
self.bn = BatchNormalization()
self.out = Dense(classes)
self.softmax = Activation("sigmoid")
self.sigmoid = Activation("sigmoid")
def call(self,inputs):
self.base_model.trainable = False
x = self.base_model(inputs, training=False)
x = self.flatten(x)
x = self.dense(x)
x = self.bn(x)
x = self.out(x)
if self.classes > 1:
x = self.softmax(x)
else:
x = self.sigmoid(x)
return x