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
|
||||
|
||||
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)
|
||||
|
|
Reference in a new issue