Spaces:
Sleeping
Sleeping
import cv2 | |
import numpy as np | |
from tensorflow.keras.applications.resnet import ResNet152 | |
from tensorflow.keras.layers import AveragePooling2D, Flatten | |
from tensorflow.keras.models import Model | |
class FeatureExtractor: | |
def __init__(self, img_shape, channels, seq_length): | |
self.seq_length = seq_length | |
self.height = img_shape[0] | |
self.width = img_shape[1] | |
self.channels = channels | |
self.base_model = ResNet152(include_top=False, input_shape=(224, 224, 3), weights='imagenet') | |
for layer in self.base_model.layers: | |
layer.trainable = False | |
self.op = self.base_model.output | |
self.x_model = AveragePooling2D((7, 7), name='avg_pool')(self.op) | |
self.x_model = Flatten()(self.x_model) | |
self.model = Model(self.base_model.input, self.x_model) | |
def extract_feature(self, frames_buffer): | |
x_op = np.zeros((2048, self.seq_length)) | |
for i in range(len(frames_buffer)): | |
x_t = frames_buffer[i] | |
x_t = cv2.resize(x_t, (224, 224)) | |
x_t = np.expand_dims(x_t, axis=0) | |
x = self.model.predict(x_t) | |
x_op[:, i] = x.flatten() | |
return x_op |