Spaces:
Runtime error
Runtime error
File size: 6,122 Bytes
1cf3825 | 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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | import tensorflow as tf
from tensorflow.keras import layers
def build_cnn_baseline(input_shape=(224, 224, 3), dropout_rate=0.3):
inputs = tf.keras.Input(shape=input_shape, name='image_input')
x = layers.Rescaling(1.0 / 255)(inputs)
x = layers.Conv2D(32, 3, activation='relu', padding='same', name='conv_block_1')(x)
x = layers.MaxPooling2D(name='pool_block_1')(x)
x = layers.Conv2D(64, 3, activation='relu', padding='same', name='conv_block_2')(x)
x = layers.MaxPooling2D(name='pool_block_2')(x)
x = layers.Conv2D(128, 3, activation='relu', padding='same', name='conv_block_3')(x)
x = layers.MaxPooling2D(name='pool_block_3')(x)
x = layers.Dropout(dropout_rate)(x)
x = layers.Flatten()(x)
x = layers.Dense(128, activation='relu')(x)
x = layers.Dropout(dropout_rate)(x)
outputs = layers.Dense(1, activation='sigmoid', name='output')(x)
return tf.keras.Model(inputs, outputs, name='cnn_baseline')
def build_transfer_model(
input_shape=(224, 224, 3),
dropout_rate=0.3,
base_model_name='resnet50',
weights='imagenet',
fine_tune=False,
fine_tune_at=None,
):
inputs = tf.keras.Input(shape=input_shape, name='image_input')
if base_model_name.lower() == 'resnet50':
x = tf.keras.applications.resnet50.preprocess_input(inputs)
base_model = tf.keras.applications.ResNet50(
include_top=False,
weights=weights,
input_shape=input_shape,
pooling='avg',
name='resnet_base',
)
elif base_model_name.lower() == 'vgg16':
x = tf.keras.applications.vgg16.preprocess_input(inputs)
base_model = tf.keras.applications.VGG16(
include_top=False,
weights=weights,
input_shape=input_shape,
pooling='avg',
name='vgg_base',
)
else:
raise ValueError('Unsupported base_model_name. Use resnet50 or vgg16.')
base_model.trainable = fine_tune
if fine_tune and fine_tune_at is not None:
for layer in base_model.layers[:fine_tune_at]:
layer.trainable = False
for layer in base_model.layers:
if isinstance(layer, layers.BatchNormalization):
layer.trainable = False
x = base_model(x, training=False)
x = layers.Dropout(dropout_rate)(x)
x = layers.Dense(256, activation='relu')(x)
x = layers.Dropout(dropout_rate)(x)
outputs = layers.Dense(1, activation='sigmoid', name='output')(x)
return tf.keras.Model(inputs, outputs, name='transfer_model')
class PatchEncoder(layers.Layer):
def __init__(self, num_patches, projection_dim):
super().__init__()
self.num_patches = num_patches
self.projection = layers.Dense(projection_dim)
self.position_embedding = layers.Embedding(input_dim=num_patches, output_dim=projection_dim)
def call(self, patch):
positions = tf.range(start=0, limit=self.num_patches, delta=1)
encoded = self.projection(patch) + self.position_embedding(positions)
return encoded
def transformer_block(x, num_heads, projection_dim, mlp_dim, dropout_rate, block_id=None):
x1 = layers.LayerNormalization(epsilon=1e-6, name=f'vit_norm_{block_id}_1')(x)
attention_layer = layers.MultiHeadAttention(
num_heads=num_heads,
key_dim=projection_dim,
dropout=dropout_rate,
name=f'vit_attention_{block_id}',
)
attention_output = attention_layer(x1, x1)
x2 = layers.Add()([attention_output, x])
x3 = layers.LayerNormalization(epsilon=1e-6, name=f'vit_norm_{block_id}_2')(x2)
x3 = layers.Dense(mlp_dim, activation='gelu')(x3)
x3 = layers.Dropout(dropout_rate)(x3)
x3 = layers.Dense(projection_dim)(x3)
x3 = layers.Dropout(dropout_rate)(x3)
x4 = layers.Add()([x3, x2])
return x4
def build_vit_classifier(
input_shape=(224, 224, 3),
patch_size=16,
num_layers=4,
num_heads=4,
projection_dim=128,
mlp_dim=256,
dropout_rate=0.1,
weights='imagenet',
):
inputs = tf.keras.Input(shape=input_shape, name='image_input')
x = tf.keras.applications.resnet50.preprocess_input(inputs)
backbone = tf.keras.applications.ResNet50(
include_top=False,
weights=weights,
input_shape=input_shape,
pooling=None,
name='vit_hybrid_resnet_base',
)
backbone.trainable = False
x = backbone(x, training=False)
x = layers.Conv2D(projection_dim, 1, padding='same', name='hybrid_patch_projection')(x)
num_patches = (input_shape[0] // 32) * (input_shape[1] // 32)
patches = layers.Reshape((num_patches, projection_dim), name='hybrid_patch_tokens')(x)
x = PatchEncoder(num_patches, projection_dim)(patches)
for i in range(num_layers):
x = transformer_block(x, num_heads, projection_dim, mlp_dim, dropout_rate, block_id=i)
x = layers.LayerNormalization(epsilon=1e-6)(x)
x = layers.GlobalAveragePooling1D()(x)
x = layers.Dropout(dropout_rate)(x)
x = layers.Dense(128, activation='relu')(x)
x = layers.Dropout(dropout_rate)(x)
outputs = layers.Dense(1, activation='sigmoid', name='output')(x)
return tf.keras.Model(inputs=inputs, outputs=outputs, name='vit_classifier')
def get_model(
model_name,
input_shape=(224, 224, 3),
transfer_weights='imagenet',
fine_tune_transfer=False,
transfer_fine_tune_at=None,
):
model_name = model_name.lower()
if model_name == 'cnn':
return build_cnn_baseline(input_shape=input_shape)
if model_name == 'transfer':
return build_transfer_model(
input_shape=input_shape,
weights=transfer_weights,
fine_tune=fine_tune_transfer,
fine_tune_at=transfer_fine_tune_at,
)
if model_name == 'vit':
return build_vit_classifier(input_shape=input_shape, weights=transfer_weights)
raise ValueError('Unknown model_name. Use cnn, transfer, or vit.')
|