49 lines
1.8 KiB
Python
49 lines
1.8 KiB
Python
# video_compression_model.py
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
|
|
from global_train import LOGGER
|
|
|
|
PRESET_SPEED_CATEGORIES = ["ultrafast", "superfast", "veryfast", "faster", "fast", "medium", "slow", "slower", "veryslow"]
|
|
NUM_PRESET_SPEEDS = len(PRESET_SPEED_CATEGORIES)
|
|
NUM_CHANNELS = 3
|
|
WIDTH = 640
|
|
HEIGHT = 360
|
|
|
|
#from tensorflow.keras.mixed_precision import Policy
|
|
|
|
#policy = Policy('mixed_float16')
|
|
#tf.keras.mixed_precision.set_global_policy(policy)
|
|
|
|
|
|
|
|
class VideoCompressionModel(tf.keras.Model):
|
|
def __init__(self):
|
|
super(VideoCompressionModel, self).__init__()
|
|
LOGGER.debug("Initializing VideoCompressionModel.")
|
|
|
|
# Add an additional channel for the histogram features
|
|
input_shape_with_histogram = (HEIGHT, WIDTH, 2) # 1 channel for edges, 1 for histogram
|
|
|
|
self.encoder = tf.keras.Sequential([
|
|
tf.keras.layers.InputLayer(input_shape=input_shape_with_histogram),
|
|
tf.keras.layers.Conv2D(64, (3, 3), activation='relu', padding='same'),
|
|
tf.keras.layers.MaxPooling2D((2, 2), padding='same'),
|
|
tf.keras.layers.Conv2D(32, (3, 3), activation='relu', padding='same'),
|
|
tf.keras.layers.MaxPooling2D((2, 2), padding='same')
|
|
])
|
|
|
|
self.decoder = tf.keras.Sequential([
|
|
tf.keras.layers.Conv2DTranspose(32, (3, 3), activation='relu', padding='same'),
|
|
tf.keras.layers.UpSampling2D((2, 2)),
|
|
tf.keras.layers.Conv2DTranspose(64, (3, 3), activation='relu', padding='same'),
|
|
tf.keras.layers.UpSampling2D((2, 2)),
|
|
tf.keras.layers.Conv2DTranspose(1, (3, 3), activation='sigmoid', padding='same')
|
|
])
|
|
|
|
def call(self, inputs):
|
|
encoded = self.encoder(inputs)
|
|
decoded = self.decoder(encoded)
|
|
return decoded
|