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

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)