File size: 3,944 Bytes
9339e0c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | import numpy as np
from sklearn.preprocessing import LabelEncoder
import torch
import os
from torchvision import models
import time
import json
def extract_features(samples, transform, dataset):
print("Extracting features using ResNet50...")
# Load the pre-trained ResNet50 model
start_time = time.time()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT).to(device)
resnet.eval()
# Remove the last layer of the ResNet50 model to obtain the feature extractor
resnet_feat = torch.nn.Sequential(*list(resnet.children())[:-1]).to(device)
resnet_feat.eval()
processed_samples = []
for frames, label in samples:
transformed_frames = [transform(frame).to(device) for frame in frames]
frames_tensor = torch.stack(transformed_frames, dim=0).to(device)
with torch.no_grad():
features_tensor = resnet_feat(frames_tensor) # Shape: (T, 2048, 1, 1)
features = torch.flatten(features_tensor, start_dim=1).cpu().numpy()
processed_samples.append((features, label))
end_time = time.time()
print(f"Feature extraction completed in {end_time - start_time:.2f} seconds")
os.makedirs(f'./benchmarks/{dataset}', exist_ok=True)
# Save the feature extraction time in file .txt
with open(f'./benchmarks/{dataset}/benchmark.txt', 'w') as f:
f.write(f"Feature extraction time: {end_time - start_time:.2f} seconds\n")
return processed_samples
def splittingData(samples, dataset):
# Shuffle the samples
np.random.shuffle(samples)
# Split the samples into training and testing sets (80% training, 20% testing)
split_idx = int(0.8 * len(samples))
train_samples = samples[:split_idx]
test_samples = samples[split_idx:]
# Split the training sets into validation and training sets (80% training, 20% validation)
validation_split_idx = int(0.8 * len(train_samples))
train_samples, val_samples = train_samples[:validation_split_idx], train_samples[validation_split_idx:]
# Separate features and labels for training, validation, and testing sets
train_features, train_labels = zip(*train_samples)
val_features, val_labels = zip(*val_samples)
test_features, test_labels = zip(*test_samples)
# Convert the labels to numerical labels using a LabelEncoder
le = LabelEncoder()
train_labels = le.fit_transform(train_labels)
val_labels = le.transform(val_labels)
test_labels = le.transform(test_labels)
train_features = np.array(train_features)
val_features = np.array(val_features)
test_features = np.array(test_features)
# Print the shapes of the features and labels arrays
print("Train Features shape:", train_features.shape)
print("Train Labels shape:", train_labels.shape)
print("Validation Features shape:", val_features.shape)
print("Validation Labels shape:", val_labels.shape)
print("Test Features shape:", test_features.shape)
print("Test Labels shape:", test_labels.shape)
os.makedirs(f'./features/{dataset}', exist_ok=True)
# Save the features and labels to numpy arrays
np.save(f'./features/{dataset}/train_features.npy', train_features)
np.save(f'./features/{dataset}/train_labels.npy', train_labels)
np.save(f'./features/{dataset}/val_features.npy', val_features)
np.save(f'./features/{dataset}/val_labels.npy', val_labels)
np.save(f'./features/{dataset}/test_features.npy', test_features)
np.save(f'./features/{dataset}/test_labels.npy', test_labels)
idx2label = {i: label for i, label in enumerate(le.classes_)}
with open(f'./features/{dataset}/label_map_idx2label.json', 'w') as f:
json.dump(idx2label, f, indent=4)
# Save the LabelEncoder for later use
return le
|