74 lines
2 KiB
Python
74 lines
2 KiB
Python
# featureExtraction.py
|
|
|
|
import cv2
|
|
import numpy as np
|
|
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, LOGGER, NUM_PRESET_SPEEDS, WIDTH
|
|
|
|
def scale_crf(crf):
|
|
return crf / 51
|
|
|
|
|
|
def scale_speed_preset(speed):
|
|
return speed / NUM_PRESET_SPEEDS
|
|
|
|
|
|
def psnr(y_true, y_pred):
|
|
#LOGGER.info(f"[psnr function] y_true: {y_true.shape}, y_pred: {y_pred.shape}")
|
|
max_pixel = 1.0
|
|
mse = K.mean(K.square(y_pred - y_true))
|
|
return 20.0 * K.log(max_pixel / K.sqrt(mse)) / 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 combined_loss(y_true, y_pred):
|
|
return -combined(y_true, y_pred) # The goal is to maximize the combined value
|
|
|
|
|
|
def detect_noise(image, threshold=15):
|
|
# Convert to grayscale if it's a color image
|
|
if len(image.shape) == 3:
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
|
|
# Compute the standard deviation
|
|
std_dev = np.std(image)
|
|
|
|
# If the standard deviation is higher than a threshold, it might be considered noisy
|
|
return std_dev > threshold
|
|
|
|
|
|
def frame_difference(frame1, frame2):
|
|
# Ensure both frames are of the same size and type
|
|
if frame1.shape != frame2.shape:
|
|
raise ValueError("Frames must have the same dimensions and number of channels")
|
|
|
|
# Calculate the absolute difference between the frames
|
|
difference = cv2.absdiff(frame1, frame2)
|
|
|
|
return difference
|
|
|
|
|
|
def preprocess_frame(frame, resize=True, scale=True):
|
|
|
|
# Check frame dimensions and resize if necessary
|
|
if resize and frame.shape[:2] != (HEIGHT, WIDTH):
|
|
frame = cv2.resize(frame, (WIDTH, HEIGHT), interpolation=cv2.INTER_LINEAR)
|
|
|
|
if scale:
|
|
# Scale frame to [0, 1]
|
|
frame = frame / 255.0
|
|
|
|
return frame
|