Spaces:
Sleeping
Sleeping
import torch | |
from torch.utils.data import DataLoader, random_split | |
from datasets import load_dataset | |
from src.dataset import HumanActionDataset | |
from utils.preprocessing import get_transforms | |
from models.resnet_model import ResNet18 | |
import pandas as pd | |
def get_val_loader(batch_size=32): | |
ds = load_dataset("Bingsu/Human_Action_Recognition") | |
full_dataset = HumanActionDataset(ds['train'], transform=get_transforms()) | |
train_size = int(0.8 * len(full_dataset)) | |
val_size = len(full_dataset) - train_size | |
_, val_dataset = random_split(full_dataset, [train_size, val_size]) | |
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False) | |
return val_loader | |
def get_test_loader(batch_size=32): | |
ds = load_dataset("Bingsu/Human_Action_Recognition") | |
test_dataset = HumanActionDataset(ds['test'], transform=get_transforms()) | |
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) | |
return test_loader | |
def evaluate_model(model_path, batch_size=32): | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
print(f"Using device: {device}") | |
model = ResNet18(num_classes=15) | |
model.load_state_dict(torch.load(model_path, map_location=device)) | |
model.to(device) | |
model.eval() | |
val_loader = get_val_loader(batch_size) | |
correct = 0 | |
total = 0 | |
with torch.no_grad(): | |
for images, labels in val_loader: | |
images = images.to(device) | |
labels = labels.to(device) | |
outputs = model(images) | |
_, predicted = torch.max(outputs, 1) | |
total += labels.size(0) | |
correct += (predicted == labels).sum().item() | |
accuracy = correct / total | |
print(f"Validation Accuracy: {accuracy:.4f} ({correct}/{total})") | |
def predict_test(model_path, batch_size=32, output_csv="test_predictions.csv"): | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
print(f"Using device: {device}") | |
model = ResNet18(num_classes=15) | |
model.load_state_dict(torch.load(model_path, map_location=device)) | |
model.to(device) | |
model.eval() | |
test_loader = get_test_loader(batch_size) | |
class_names = [ | |
'calling', 'clapping', 'cycling', 'dancing', 'drinking', 'eating', 'fighting', | |
'hugging', 'laughing', 'listening_to_music', 'running', 'sitting', 'sleeping', | |
'texting', 'using_laptop' | |
] | |
predictions = [] | |
filenames = [] | |
with torch.no_grad(): | |
for batch in test_loader: | |
images = batch[0].to(device) | |
# Assuming your HumanActionDataset returns (image, label, filename) or just (image, label) | |
# If filenames are needed and your dataset doesn't provide, you can omit or modify accordingly. | |
outputs = model(images) | |
_, predicted = torch.max(outputs, 1) | |
predictions.extend(predicted.cpu().tolist()) | |
# If you have filenames: | |
# filenames.extend(batch[2]) # Uncomment if your dataset yields filenames | |
# If you don't have filenames, just save predictions indexed by order | |
df = pd.DataFrame({ | |
"id": list(range(len(predictions))), | |
"label": [class_names[p] for p in predictions] | |
}) | |
df.to_csv(output_csv, index=False) | |
print(f"Saved test predictions to {output_csv}") | |
if __name__ == "__main__": | |
model_path = "models/best_model.pth" | |
evaluate_model(model_path=model_path, batch_size=32) | |
predict_test(model_path=model_path, batch_size=32, output_csv="test_predictions.csv") | |