test
This commit is contained in:
parent
8c5001166d
commit
5085c87300
3 changed files with 96 additions and 173 deletions
103
DeepEncode.py
103
DeepEncode.py
|
|
@ -1,91 +1,68 @@
|
|||
import tensorflow as tf
|
||||
import numpy as np
|
||||
import cv2
|
||||
from video_compression_model import NUM_FRAMES, PRESET_SPEED_CATEGORIES, VideoCompressionModel
|
||||
from video_compression_model import VideoCompressionModel
|
||||
|
||||
# Constants
|
||||
MAX_FRAMES = 24
|
||||
CHUNK_SIZE = 24 # Adjust based on available memory and video resolution
|
||||
COMPRESSED_VIDEO_FILE = 'compressed_video.mkv'
|
||||
COMPRESSED_VIDEO_FILE = 'compressed_video.mp4'
|
||||
MAX_FRAMES = 24 # Limit the number of frames processed
|
||||
|
||||
# Load the trained model
|
||||
model = tf.keras.models.load_model('models/model.keras', custom_objects={'VideoCompressionModel': VideoCompressionModel})
|
||||
|
||||
# Step 2: Load the trained model
|
||||
model = tf.keras.models.load_model('models/model_differencing.keras', custom_objects={'VideoCompressionModel': VideoCompressionModel})
|
||||
|
||||
# Step 3: Load the uncompressed video
|
||||
# Load the uncompressed video
|
||||
UNCOMPRESSED_VIDEO_FILE = 'test_data/training_video.mkv'
|
||||
|
||||
def load_frames_from_video(video_file, start_frame=0, num_frames=CHUNK_SIZE):
|
||||
def load_frame_from_video(video_file, frame_num):
|
||||
cap = cv2.VideoCapture(video_file)
|
||||
frames = []
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
|
||||
|
||||
for _ in range(num_frames):
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 # Normalize and convert to float32
|
||||
frames.append(frame)
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, 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
|
||||
cap.release()
|
||||
return frames
|
||||
|
||||
def predict_in_chunks(uncompressed_frames, model, crf_values, preset_speed_values):
|
||||
num_sequences = len(uncompressed_frames) - NUM_FRAMES + 1
|
||||
compressed_frames = []
|
||||
|
||||
#for frame in uncompressed_frames:
|
||||
# cv2.imshow("frame", frame)
|
||||
# cv2.waitKey(50)
|
||||
#display_frame = np.clip(frame * 255.0, 0, 255).astype(np.uint8)
|
||||
#cv2.imshow("uncomp", display_frame)
|
||||
#cv2.waitKey(0) # Add this line to hold the display window until a key is pressed
|
||||
|
||||
for start in range(0, num_sequences, CHUNK_SIZE):
|
||||
end = min(start + CHUNK_SIZE, num_sequences)
|
||||
frame_chunk = uncompressed_frames[start:end + NUM_FRAMES - 1]
|
||||
crf_chunk = crf_values[start:end]
|
||||
speed_chunk = preset_speed_values[start:end]
|
||||
|
||||
frame_sequences = []
|
||||
for i in range(len(frame_chunk) - NUM_FRAMES + 1):
|
||||
sequence = frame_chunk[i:i + NUM_FRAMES]
|
||||
frame_sequences.append(sequence)
|
||||
|
||||
frame_sequences = np.array(frame_sequences)
|
||||
|
||||
compressed_chunk = model.predict({"frames": frame_sequences, "crf": crf_chunk, "preset_speed": speed_chunk})
|
||||
compressed_frames.extend(compressed_chunk)
|
||||
|
||||
return compressed_frames
|
||||
|
||||
def save_frames_chunk(frames, video_writer):
|
||||
for frame in frames:
|
||||
frame = np.clip(frame * 255.0, 0, 255).astype(np.uint8)
|
||||
frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
|
||||
video_writer.write(frame)
|
||||
|
||||
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])
|
||||
|
||||
compressed_frame = model.predict({
|
||||
"frame": np.array([uncompressed_frame]),
|
||||
"crf": crf_array,
|
||||
"preset_speed": preset_speed_array
|
||||
})
|
||||
return compressed_frame[0]
|
||||
|
||||
cap = cv2.VideoCapture(UNCOMPRESSED_VIDEO_FILE)
|
||||
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
cap.release()
|
||||
|
||||
if MAX_FRAMES != 0 and total_frames > MAX_FRAMES:
|
||||
total_frames = MAX_FRAMES
|
||||
|
||||
cap.release()
|
||||
crf_value = 25.0 # Example CRF value
|
||||
preset_speed_value = 2 # Index for "fast" in our defined list
|
||||
|
||||
crf_values = np.full((CHUNK_SIZE + NUM_FRAMES - 1, 1), 25, dtype=np.float32) # Chunk size + look-ahead frames
|
||||
preset_speed_index = PRESET_SPEED_CATEGORIES.index("fast")
|
||||
preset_speed_values = np.full((CHUNK_SIZE + NUM_FRAMES - 1, 1), preset_speed_index, dtype=np.float32)
|
||||
height, width = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
fourcc = cv2.VideoWriter_fourcc(*'H264')
|
||||
out = cv2.VideoWriter(COMPRESSED_VIDEO_FILE, fourcc, 24.0, (width, height))
|
||||
|
||||
out = None # Video writer instance
|
||||
for i in range(0, total_frames, CHUNK_SIZE):
|
||||
uncompressed_frames_chunk = load_frames_from_video(UNCOMPRESSED_VIDEO_FILE, start_frame=i)
|
||||
compressed_frames_chunk = predict_in_chunks(uncompressed_frames_chunk, model, crf_values, preset_speed_values)
|
||||
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_value, preset_speed_value)
|
||||
|
||||
# Initialize video writer if it's the first chunk
|
||||
if out is None:
|
||||
height, width = compressed_frames_chunk[0].shape[:2]
|
||||
fourcc = cv2.VideoWriter_fourcc(*'XVID')
|
||||
out = cv2.VideoWriter(COMPRESSED_VIDEO_FILE, fourcc, 24.0, (width, height))
|
||||
|
||||
save_frames_chunk(compressed_frames_chunk, out)
|
||||
compressed_frame = np.clip(compressed_frame * 255.0, 0, 255).astype(np.uint8)
|
||||
compressed_frame = cv2.cvtColor(compressed_frame, cv2.COLOR_RGB2BGR)
|
||||
out.write(compressed_frame)
|
||||
cv2.imshow("output", compressed_frame)
|
||||
|
||||
out.release()
|
||||
print("Compression completed.")
|
||||
|
|
|
|||
Reference in a new issue