jordon/DeepEncode
Archived
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update

Browse source
This commit is contained in:
Jordon Brooks 2023-08-17 00:00:19 +01:00
parent 6a2449c5cd
commit 046b78cb87
3 changed files with 23 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

11
DeepEncode.py
View file

@ -2,7 +2,7 @@
import os
from featureExtraction import preprocess_frame, psnr
from featureExtraction import preprocess_frame, psnr, scale_crf, scale_speed_preset
from globalVars import PRESET_SPEED_CATEGORIES
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
@ -10,12 +10,12 @@ os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
import tensorflow as tf
import numpy as np
import cv2
from video_compression_model import VideoCompressionModel
from video_compression_model import VideoCompressionModel, combine_batch
# Constants
COMPRESSED_VIDEO_FILE = 'compressed_video.avi'
MAX_FRAMES = 0 # Limit the number of frames processed
CRF = 51
CRF = 0
SPEED = PRESET_SPEED_CATEGORIES.index("ultrafast")
# Load the trained model
@ -38,10 +38,13 @@ def predict_frame(uncompressed_frame):
#display_frame = np.clip(cv2.cvtColor(uncompressed_frame, cv2.COLOR_BGR2RGB) * 255.0, 0, 255).astype(np.uint8)
#cv2.imshow("uncomp", uncompressed_frame)
scaled_crf = scale_crf(CRF)
scaled_speed = scale_speed_preset(SPEED)
frame = preprocess_frame(uncompressed_frame, CRF, SPEED)
frame = combine_batch(uncompressed_frame, scaled_crf, scaled_speed)
compressed_frame = MODEL.predict([np.expand_dims(frame, axis=0)])[0]
compressed_frame = compressed_frame[:, :, :3] # Keep only the first 3 channels (BGR)
display_frame = np.clip(compressed_frame * 255.0, 0, 255).astype(np.uint8)

6
featureExtraction.py
View file

@ -10,6 +10,12 @@ from tensorflow.keras import backend as K
from globalVars import HEIGHT, NUM_PRESET_SPEEDS, WIDTH
def scale_crf(crf):
return crf / 51
def scale_speed_preset(speed):
return speed / NUM_PRESET_SPEEDS
def extract_edge_features(frame):
"""

25
video_compression_model.py
View file

@ -4,7 +4,7 @@ import os
import cv2
import numpy as np
import tensorflow as tf
from featureExtraction import preprocess_frame
from featureExtraction import preprocess_frame, scale_crf, scale_speed_preset
from globalVars import HEIGHT, LOGGER, NUM_COLOUR_CHANNELS, NUM_PRESET_SPEEDS, PRESET_SPEED_CATEGORIES, WIDTH
@ -13,10 +13,13 @@ from globalVars import HEIGHT, LOGGER, NUM_COLOUR_CHANNELS, NUM_PRESET_SPEEDS, P
#policy = Policy('mixed_float16')
#tf.keras.mixed_precision.set_global_policy(policy)
def combine_batch(frame, crf_array, speed_array):
def combine_batch(frame, crf, speed):
# Preprocess the compressed frame (target)
processed_frame = preprocess_frame(frame)
crf_array = np.full((HEIGHT, WIDTH, 1), crf) # Note the added dimension
speed_array = np.full((HEIGHT, WIDTH, 1), speed) # Note the added dimension
# Combine the frames with the CRF and SPEED images
combined = np.concatenate([processed_frame, crf_array, speed_array], axis=-1)
@ -32,16 +35,8 @@ def data_generator(videos, batch_size):
video_path = os.path.join(base_dir, video_details["compressed_video_file"])
uncompressed_video_path = os.path.join(base_dir, video_details["original_video_file"])
CRF = video_details["crf"] / 51
SPEED = PRESET_SPEED_CATEGORIES.index(video_details["preset_speed"]) / NUM_PRESET_SPEEDS
# Create images with the CRF and SPEED values, filling extra channels
compressed_crf_array = np.full((HEIGHT, WIDTH, 1), CRF) # Note the added dimension
compressed_speed_array = np.full((HEIGHT, WIDTH, 1), SPEED) # Note the added dimension
# Create images with the CRF and SPEED values, filling extra channels
uncompressed_crf_array = np.full((HEIGHT, WIDTH, 1), 0) # Note the added dimension
uncompressed_speed_array = np.full((HEIGHT, WIDTH, 1), PRESET_SPEED_CATEGORIES.index("veryslow") / NUM_PRESET_SPEEDS) # Note the added dimension
CRF = scale_crf(video_details["crf"])
SPEED = scale_speed_preset(PRESET_SPEED_CATEGORIES.index(video_details["preset_speed"]))
# Open the video files
cap_compressed = cv2.VideoCapture(video_path)
@ -58,9 +53,9 @@ def data_generator(videos, batch_size):
if not ret_compressed or not ret_uncompressed:
break
compressed_combined = combine_batch(compressed_frame, compressed_crf_array, compressed_speed_array)
compressed_combined = combine_batch(compressed_frame, CRF, SPEED)
uncompressed_combined = combine_batch(uncompressed_frame, uncompressed_crf_array, uncompressed_speed_array)
uncompressed_combined = combine_batch(uncompressed_frame, 0, scale_speed_preset(PRESET_SPEED_CATEGORIES.index("veryslow")))
# Append processed frames to batches
compressed_frame_batch.append(compressed_combined)
@ -103,7 +98,7 @@ class VideoCompressionModel(tf.keras.Model):
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(NUM_COLOUR_CHANNELS + 2, (3, 3), activation='sigmoid', padding='same')
tf.keras.layers.Conv2DTranspose(NUM_COLOUR_CHANNELS, (3, 3), activation='sigmoid', padding='same')
])
def call(self, inputs):