61 lines
2.8 KiB
Python
61 lines
2.8 KiB
Python
import tensorflow as tf
|
|
|
|
PRESET_SPEED_CATEGORIES = ["ultrafast", "superfast", "veryfast", "faster", "fast", "medium", "slow", "slower", "veryslow"]
|
|
NUM_PRESET_SPEEDS = len(PRESET_SPEED_CATEGORIES)
|
|
NUM_CHANNELS = 3 # Number of color channels in the video frames (RGB images have 3 channels)
|
|
|
|
#policy = tf.keras.mixed_precision.Policy('mixed_float16')
|
|
#tf.keras.mixed_precision.set_global_policy(policy)
|
|
|
|
class VideoCompressionModel(tf.keras.Model):
|
|
def __init__(self, NUM_CHANNELS=3, NUM_FRAMES=5, regularization_factor=1e-4):
|
|
super(VideoCompressionModel, self).__init__()
|
|
|
|
self.NUM_CHANNELS = NUM_CHANNELS
|
|
|
|
# Regularization
|
|
self.regularizer = tf.keras.regularizers.l2(regularization_factor)
|
|
|
|
# Embedding layer for preset_speed
|
|
self.preset_embedding = tf.keras.layers.Embedding(NUM_PRESET_SPEEDS, 16, embeddings_regularizer=self.regularizer)
|
|
|
|
# Encoder layers
|
|
self.encoder = tf.keras.Sequential([
|
|
tf.keras.layers.ZeroPadding2D(padding=((1, 1), (1, 1))), # Padding to preserve spatial dimensions
|
|
tf.keras.layers.Conv2D(32, (3, 3), activation='relu', padding='same', kernel_regularizer=self.regularizer),
|
|
tf.keras.layers.MaxPooling2D((2, 2)),
|
|
# Add more encoder layers as needed
|
|
])
|
|
|
|
# Decoder layers
|
|
self.decoder = tf.keras.Sequential([
|
|
tf.keras.layers.Conv2DTranspose(32, (3, 3), activation='relu', padding='same', kernel_regularizer=self.regularizer),
|
|
tf.keras.layers.UpSampling2D((2, 2)),
|
|
# Add more decoder layers as needed
|
|
tf.keras.layers.Conv2D(NUM_CHANNELS, (3, 3), activation='sigmoid', padding='same', kernel_regularizer=self.regularizer), # Output layer for video frames
|
|
tf.keras.layers.Cropping2D(cropping=((1, 1), (1, 1))) # Adjust cropping to ensure dimensions match
|
|
|
|
])
|
|
|
|
def call(self, inputs):
|
|
frame = inputs["frame"]
|
|
crf = tf.expand_dims(inputs["crf"], -1)
|
|
preset_speed = inputs["preset_speed"]
|
|
|
|
# Convert preset_speed to embeddings
|
|
preset_embedding = self.preset_embedding(preset_speed)
|
|
preset_embedding = tf.keras.layers.Flatten()(preset_embedding)
|
|
|
|
# Concatenate crf and preset_embedding to frames
|
|
frame_shape = tf.shape(frame)
|
|
repeated_crf = tf.tile(tf.reshape(crf, (-1, 1, 1, 1)), [1, frame_shape[1], frame_shape[2], 1])
|
|
repeated_preset = tf.tile(tf.reshape(preset_embedding, (-1, 1, 1, 16)), [1, frame_shape[1], frame_shape[2], 1])
|
|
|
|
frame = tf.concat([tf.cast(frame, tf.float32), repeated_crf, repeated_preset], axis=-1)
|
|
|
|
# Encoding the frame
|
|
compressed_representation = self.encoder(frame)
|
|
|
|
# Decoding to generate compressed frame
|
|
reconstructed_frame = self.decoder(compressed_representation)
|
|
return reconstructed_frame
|