Update
This commit is contained in:
parent
54fa90247a
commit
15d8e57da5
4 changed files with 56 additions and 12 deletions
|
@ -2,7 +2,7 @@
|
|||
|
||||
import os
|
||||
|
||||
from featureExtraction import preprocess_frame
|
||||
from featureExtraction import preprocess_frame, psnr
|
||||
from globalVars import PRESET_SPEED_CATEGORIES
|
||||
|
||||
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
|
||||
|
@ -16,10 +16,10 @@ from video_compression_model import VideoCompressionModel
|
|||
COMPRESSED_VIDEO_FILE = 'compressed_video.avi'
|
||||
MAX_FRAMES = 0 # Limit the number of frames processed
|
||||
CRF = 51
|
||||
SPEED = PRESET_SPEED_CATEGORIES.index("veryslow")
|
||||
SPEED = PRESET_SPEED_CATEGORIES.index("ultrafast")
|
||||
|
||||
# Load the trained model
|
||||
MODEL = tf.keras.models.load_model('models/model.tf', custom_objects={'VideoCompressionModel': VideoCompressionModel})
|
||||
MODEL = tf.keras.models.load_model('models/model.tf', custom_objects={'VideoCompressionModel': VideoCompressionModel, 'psnr': psnr})
|
||||
|
||||
# Load the uncompressed video
|
||||
UNCOMPRESSED_VIDEO_FILE = 'test_data/training_video.mkv'
|
||||
|
@ -36,8 +36,8 @@ def load_frame_from_video(video_file, frame_num):
|
|||
|
||||
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)
|
||||
#display_frame = np.clip(cv2.cvtColor(uncompressed_frame, cv2.COLOR_BGR2RGB) * 255.0, 0, 255).astype(np.uint8)
|
||||
#cv2.imshow("uncomp", uncompressed_frame)
|
||||
|
||||
frame = preprocess_frame(uncompressed_frame, CRF, SPEED)
|
||||
|
||||
|
|
Reference in a new issue