updated
This commit is contained in:
parent
60c6c97071
commit
93ccce5ec1
6 changed files with 46 additions and 34 deletions
|
@ -3,14 +3,14 @@
|
|||
import tensorflow as tf
|
||||
import numpy as np
|
||||
import cv2
|
||||
from video_compression_model import VideoCompressionModel
|
||||
from video_compression_model import PRESET_SPEED_CATEGORIES, VideoCompressionModel
|
||||
|
||||
# Constants
|
||||
CHUNK_SIZE = 10 # Adjust based on available memory and video resolution
|
||||
CHUNK_SIZE = 24 # Adjust based on available memory and video resolution
|
||||
COMPRESSED_VIDEO_FILE = 'compressed_video.avi'
|
||||
MAX_FRAMES = 0 # Limit the number of frames processed
|
||||
CRF = 25.0 # Example CRF value
|
||||
PRESET_SPEED = 4 # Index for "fast" in our defined list
|
||||
CRF = 24.0 # Example CRF value
|
||||
PRESET_SPEED = "veryslow" # Index for "fast" in our defined list
|
||||
|
||||
# Load the trained model
|
||||
model = tf.keras.models.load_model('models/model.tf', custom_objects={'VideoCompressionModel': VideoCompressionModel})
|
||||
|
@ -42,6 +42,7 @@ def predict_frame(uncompressed_frame, model, crf_value, preset_speed_value):
|
|||
|
||||
compressed_frame = model.predict({
|
||||
"compressed_frame": uncompressed_frame,
|
||||
"uncompressed_frame": uncompressed_frame,
|
||||
"crf": crf_array,
|
||||
"preset_speed": preset_speed_array
|
||||
})
|
||||
|
@ -49,7 +50,7 @@ def predict_frame(uncompressed_frame, model, crf_value, preset_speed_value):
|
|||
display_frame = np.clip(cv2.cvtColor(compressed_frame[0], cv2.COLOR_BGR2RGB) * 255.0, 0, 255).astype(np.uint8)
|
||||
|
||||
cv2.imshow("comp", display_frame)
|
||||
cv2.waitKey(10)
|
||||
cv2.waitKey(1)
|
||||
|
||||
return compressed_frame[0]
|
||||
|
||||
|
@ -70,7 +71,7 @@ if MAX_FRAMES != 0 and total_frames > MAX_FRAMES:
|
|||
|
||||
for i in range(total_frames):
|
||||
uncompressed_frame = load_frame_from_video(UNCOMPRESSED_VIDEO_FILE, frame_num=i)
|
||||
compressed_frame = predict_frame(uncompressed_frame, model, CRF, PRESET_SPEED)
|
||||
compressed_frame = predict_frame(uncompressed_frame, model, CRF, PRESET_SPEED_CATEGORIES.index(PRESET_SPEED))
|
||||
|
||||
compressed_frame = np.clip(compressed_frame * 255.0, 0, 255).astype(np.uint8)
|
||||
compressed_frame = cv2.cvtColor(compressed_frame, cv2.COLOR_RGB2BGR)
|
||||
|
|
Reference in a new issue