update
This commit is contained in:
parent
1d98bc84a2
commit
fde856f3ec
6 changed files with 107 additions and 109 deletions
|
@ -2,22 +2,22 @@
|
|||
|
||||
import os
|
||||
|
||||
from featureExtraction import preprocess_frame
|
||||
|
||||
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
|
||||
|
||||
import tensorflow as tf
|
||||
import numpy as np
|
||||
import cv2
|
||||
from video_compression_model import PRESET_SPEED_CATEGORIES, VideoCompressionModel
|
||||
from video_compression_model import VideoCompressionModel
|
||||
|
||||
# Constants
|
||||
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 = 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})
|
||||
MODEL = tf.keras.models.load_model('models/model.tf', custom_objects={'VideoCompressionModel': VideoCompressionModel})
|
||||
|
||||
# Load the uncompressed video
|
||||
UNCOMPRESSED_VIDEO_FILE = 'test_data/training_video.mkv'
|
||||
|
@ -28,39 +28,27 @@ def load_frame_from_video(video_file, frame_num):
|
|||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
return None
|
||||
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 # Normalize and convert to float32
|
||||
#frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
cap.release()
|
||||
|
||||
return frame
|
||||
|
||||
def predict_frame(uncompressed_frame, model, crf_value, preset_speed_value):
|
||||
crf_array = np.array([crf_value])
|
||||
preset_speed_array = np.array([preset_speed_value])
|
||||
|
||||
crf_array = np.expand_dims(np.array([crf_value]), axis=-1) # Shape: (1, 1)
|
||||
preset_speed_array = np.expand_dims(np.array([preset_speed_value]), axis=-1) # Shape: (1, 1)
|
||||
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)
|
||||
cv2.waitKey(1)
|
||||
|
||||
# Expand dimensions to include batch size
|
||||
uncompressed_frame = np.expand_dims(uncompressed_frame, 0)
|
||||
|
||||
#display_frame = np.clip(cv2.cvtColor(uncompressed_frame[0], cv2.COLOR_BGR2RGB) * 255.0, 0, 255).astype(np.uint8)
|
||||
#cv2.imshow("uncomp", display_frame)
|
||||
#cv2.waitKey(0)
|
||||
combined_feature, _ = preprocess_frame(uncompressed_frame)
|
||||
|
||||
compressed_frame = model.predict({
|
||||
"compressed_frame": uncompressed_frame,
|
||||
"uncompressed_frame": uncompressed_frame,
|
||||
"crf": crf_array,
|
||||
"preset_speed": preset_speed_array
|
||||
})
|
||||
compressed_frame = MODEL.predict(np.expand_dims(combined_feature, axis=0))[0]
|
||||
|
||||
display_frame = np.clip(cv2.cvtColor(compressed_frame[0], cv2.COLOR_BGR2RGB) * 255.0, 0, 255).astype(np.uint8)
|
||||
display_frame = np.clip(cv2.cvtColor(compressed_frame, cv2.COLOR_BGR2RGB) * 255.0, 0, 255).astype(np.uint8)
|
||||
|
||||
cv2.imshow("comp", display_frame)
|
||||
cv2.waitKey(1)
|
||||
|
||||
return compressed_frame[0]
|
||||
return compressed_frame
|
||||
|
||||
cap = cv2.VideoCapture(UNCOMPRESSED_VIDEO_FILE)
|
||||
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
|
@ -79,7 +67,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_CATEGORIES.index(PRESET_SPEED))
|
||||
compressed_frame = predict_frame(uncompressed_frame)
|
||||
|
||||
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