Better

featureExtraction.py

@@ -6,6 +6,7 @@ import os
 
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
 
+import tensorflow as tf
 from tensorflow.keras import backend as K
 
 from globalVars import HEIGHT, NUM_PRESET_SPEEDS, WIDTH
@@ -49,11 +50,20 @@ def extract_histogram_features(frame, bins=64):
 
     return np.array(feature_vector)
 
+
 def psnr(y_true, y_pred):
     max_pixel = 1.0
     return 10.0 * K.log((max_pixel ** 2) / (K.mean(K.square(y_pred - y_true)))) / K.log(10.0)
 
 
+def ssim(y_true, y_pred):
+    return (tf.image.ssim(y_true, y_pred, max_val=1.0) + 1) * 50  # Normalize SSIM from [-1, 1] to [0, 100]
+
+
+def combined(y_true, y_pred):
+    return (psnr(y_true, y_pred) + ssim(y_true, y_pred)) / 2
+
+
 def preprocess_frame(frame, resize=True):
     #Preprocesses a single frame, cropping it if needed