reysarms's picture
updated
c08ab4e
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")