Black and white

DeepEncode.py

@@ -3,6 +3,7 @@
 import os
 
 from featureExtraction import preprocess_frame
+from globalVars import PRESET_SPEED_CATEGORIES
 
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
 
@@ -14,6 +15,8 @@ from video_compression_model import VideoCompressionModel
 # Constants
 COMPRESSED_VIDEO_FILE = 'compressed_video.avi'
 MAX_FRAMES = 0  # Limit the number of frames processed
+CRF = 51
+SPEED = PRESET_SPEED_CATEGORIES.index("ultrafast")
 
 # Load the trained model
 MODEL = tf.keras.models.load_model('models/model.tf', custom_objects={'VideoCompressionModel': VideoCompressionModel})
@@ -37,7 +40,7 @@ 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)
     
-    frame = preprocess_frame(uncompressed_frame)
+    frame = preprocess_frame(uncompressed_frame, CRF, SPEED)
     
     compressed_frame = MODEL.predict([np.expand_dims(frame, axis=0)])[0]
     
@@ -54,7 +57,7 @@ height, width = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), int(cap.get(cv2.CAP_PRO
 cap.release()
 
 fourcc = cv2.VideoWriter_fourcc(*'XVID')
-out = cv2.VideoWriter(COMPRESSED_VIDEO_FILE, fourcc, 24.0, (width, height))
+out = cv2.VideoWriter(COMPRESSED_VIDEO_FILE, fourcc, 24.0, (width, height), True)
 
 if not out.isOpened():
     print("Error: VideoWriter could not be opened.")