jordon/DeepEncode
Archived
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

semi-working

Browse source
This commit is contained in:
Jordon Brooks 2023-08-13 20:48:00 +01:00
parent e7af02cb4f
commit 54fa90247a
4 changed files with 38 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

10
DeepEncode.py
View file

@ -16,7 +16,7 @@ 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("ultrafast")
SPEED = PRESET_SPEED_CATEGORIES.index("veryslow")
# Load the trained model
MODEL = tf.keras.models.load_model('models/model.tf', custom_objects={'VideoCompressionModel': VideoCompressionModel})
@ -30,21 +30,21 @@ 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)
cap.release()
return frame
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)
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(cv2.cvtColor(compressed_frame, cv2.COLOR_BGR2RGB) * 255.0, 0, 255).astype(np.uint8)
display_frame = np.clip(compressed_frame * 255.0, 0, 255).astype(np.uint8)
cv2.imshow("comp", display_frame)
cv2.waitKey(1)