This commit is contained in:
Jordon Brooks 2023-08-19 01:45:58 +01:00
parent 9f34cf8074
commit 2b4664fcbb
5 changed files with 34 additions and 11 deletions

View file

@ -2,7 +2,7 @@
import os
from featureExtraction import preprocess_frame, psnr, scale_crf, scale_speed_preset
from featureExtraction import combined, preprocess_frame, psnr, scale_crf, scale_speed_preset, ssim
from globalVars import PRESET_SPEED_CATEGORIES
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
@ -14,15 +14,15 @@ 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 = 0
MAX_FRAMES = 200 # Limit the number of frames processed
CRF = 51
SPEED = PRESET_SPEED_CATEGORIES.index("ultrafast")
# Load the trained model
MODEL = tf.keras.models.load_model('models/model.tf', custom_objects={'VideoCompressionModel': VideoCompressionModel, 'psnr': psnr})
MODEL = tf.keras.models.load_model('models/model.tf', custom_objects={'VideoCompressionModel': VideoCompressionModel, 'psnr': psnr, 'ssim': ssim, 'combined': combined})
# Load the uncompressed video
UNCOMPRESSED_VIDEO_FILE = 'test_data/B4_t02.mkv'
UNCOMPRESSED_VIDEO_FILE = 'test_data/x264_crf-5_preset-veryslow.mkv'
def load_frame_from_video(video_file, frame_num):
cap = cv2.VideoCapture(video_file)