67 lines
1.7 KiB
Python
67 lines
1.7 KiB
Python
# featureExtraction.py
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import os
|
|
|
|
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
|
|
|
|
from tensorflow.keras import backend as K
|
|
|
|
from globalVars import HEIGHT, NUM_PRESET_SPEEDS, WIDTH
|
|
|
|
def scale_crf(crf):
|
|
return crf / 51
|
|
|
|
def scale_speed_preset(speed):
|
|
return speed / NUM_PRESET_SPEEDS
|
|
|
|
|
|
def extract_edge_features(frame):
|
|
"""
|
|
Extract edge features using Canny edge detection.
|
|
|
|
Args:
|
|
- frame (ndarray): Image frame.
|
|
|
|
Returns:
|
|
- ndarray: Edge feature map.
|
|
"""
|
|
edges = cv2.Canny(frame, threshold1=100, threshold2=200)
|
|
return edges.astype(np.float32) / 255.0
|
|
|
|
def extract_histogram_features(frame, bins=64):
|
|
"""
|
|
Extract histogram features from a frame with 3 channels.
|
|
|
|
Args:
|
|
- frame (ndarray): Image frame with shape (height, width, 3).
|
|
- bins (int): Number of bins for the histogram.
|
|
|
|
Returns:
|
|
- ndarray: Normalized histogram feature vector.
|
|
"""
|
|
feature_vector = []
|
|
for channel in range(3):
|
|
histogram, _ = np.histogram(frame[:,:,channel].flatten(), bins=bins, range=[0, 255])
|
|
normalized_histogram = histogram.astype(np.float32) / frame[:,:,channel].size
|
|
feature_vector.extend(normalized_histogram)
|
|
|
|
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 preprocess_frame(frame):
|
|
#Preprocesses a single frame, cropping it if needed
|
|
|
|
# Check frame dimensions and resize if necessary
|
|
if frame.shape[:2] != (HEIGHT, WIDTH):
|
|
frame = cv2.resize(frame, (WIDTH, HEIGHT), interpolation=cv2.INTER_LINEAR)
|
|
|
|
# Scale frame to [0, 1]
|
|
compressed_frame = frame / 255.0
|
|
|
|
return compressed_frame
|