Update
This commit is contained in:
parent
6a2449c5cd
commit
046b78cb87
3 changed files with 23 additions and 19 deletions
|
@ -2,7 +2,7 @@
|
||||||
|
|
||||||
import os
|
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
|
from globalVars import PRESET_SPEED_CATEGORIES
|
||||||
|
|
||||||
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
|
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 tensorflow as tf
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import cv2
|
import cv2
|
||||||
from video_compression_model import VideoCompressionModel
|
from video_compression_model import VideoCompressionModel, combine_batch
|
||||||
|
|
||||||
# Constants
|
# Constants
|
||||||
COMPRESSED_VIDEO_FILE = 'compressed_video.avi'
|
COMPRESSED_VIDEO_FILE = 'compressed_video.avi'
|
||||||
MAX_FRAMES = 0 # Limit the number of frames processed
|
MAX_FRAMES = 0 # Limit the number of frames processed
|
||||||
CRF = 51
|
CRF = 0
|
||||||
SPEED = PRESET_SPEED_CATEGORIES.index("ultrafast")
|
SPEED = PRESET_SPEED_CATEGORIES.index("ultrafast")
|
||||||
|
|
||||||
# Load the trained model
|
# 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)
|
#display_frame = np.clip(cv2.cvtColor(uncompressed_frame, cv2.COLOR_BGR2RGB) * 255.0, 0, 255).astype(np.uint8)
|
||||||
#cv2.imshow("uncomp", uncompressed_frame)
|
#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 = MODEL.predict([np.expand_dims(frame, axis=0)])[0]
|
||||||
|
|
||||||
compressed_frame = compressed_frame[:, :, :3] # Keep only the first 3 channels (BGR)
|
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)
|
display_frame = np.clip(compressed_frame * 255.0, 0, 255).astype(np.uint8)
|
||||||
|
|
|
@ -10,6 +10,12 @@ from tensorflow.keras import backend as K
|
||||||
|
|
||||||
from globalVars import HEIGHT, NUM_PRESET_SPEEDS, WIDTH
|
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):
|
def extract_edge_features(frame):
|
||||||
"""
|
"""
|
||||||
|
|
|
@ -4,7 +4,7 @@ import os
|
||||||
import cv2
|
import cv2
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import tensorflow as tf
|
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
|
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')
|
#policy = Policy('mixed_float16')
|
||||||
#tf.keras.mixed_precision.set_global_policy(policy)
|
#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)
|
# Preprocess the compressed frame (target)
|
||||||
processed_frame = preprocess_frame(frame)
|
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
|
# Combine the frames with the CRF and SPEED images
|
||||||
combined = np.concatenate([processed_frame, crf_array, speed_array], axis=-1)
|
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"])
|
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"])
|
uncompressed_video_path = os.path.join(base_dir, video_details["original_video_file"])
|
||||||
|
|
||||||
CRF = video_details["crf"] / 51
|
CRF = scale_crf(video_details["crf"])
|
||||||
SPEED = PRESET_SPEED_CATEGORIES.index(video_details["preset_speed"]) / NUM_PRESET_SPEEDS
|
SPEED = scale_speed_preset(PRESET_SPEED_CATEGORIES.index(video_details["preset_speed"]))
|
||||||
|
|
||||||
# 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
|
|
||||||
|
|
||||||
# Open the video files
|
# Open the video files
|
||||||
cap_compressed = cv2.VideoCapture(video_path)
|
cap_compressed = cv2.VideoCapture(video_path)
|
||||||
|
@ -58,9 +53,9 @@ def data_generator(videos, batch_size):
|
||||||
if not ret_compressed or not ret_uncompressed:
|
if not ret_compressed or not ret_uncompressed:
|
||||||
break
|
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
|
# Append processed frames to batches
|
||||||
compressed_frame_batch.append(compressed_combined)
|
compressed_frame_batch.append(compressed_combined)
|
||||||
|
@ -103,7 +98,7 @@ class VideoCompressionModel(tf.keras.Model):
|
||||||
tf.keras.layers.UpSampling2D((2, 2)),
|
tf.keras.layers.UpSampling2D((2, 2)),
|
||||||
tf.keras.layers.Conv2DTranspose(64, (3, 3), activation='relu', padding='same'),
|
tf.keras.layers.Conv2DTranspose(64, (3, 3), activation='relu', padding='same'),
|
||||||
tf.keras.layers.UpSampling2D((2, 2)),
|
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):
|
def call(self, inputs):
|
||||||
|
|
Reference in a new issue