testing_pytorch5_ordered.py

import os
import torch
import pandas as pd
from torchvision.transforms import transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader, Dataset
from efficientnet_pytorch import EfficientNet
from PIL import Image


class CustomTestDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.image_files = sorted(os.listdir(root_dir))  # Sort image files

    def __len__(self):
        return len(self.image_files)

    def __getitem__(self, idx):
        image_path = os.path.join(self.root_dir, self.image_files[idx])
        image = Image.open(image_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        return image, self.image_files[idx].split('.')[0]  # Return file name without extension
    
# Define paths for test dataset, class names CSV file, and saved model parameters
test_dir = 'rgb_test_crop'
class_names_file = 'category.csv'
model_path = 'newb1_train_model_augmentation_epoch_2.pth'

# Load class names from CSV file
class_names_df = pd.read_csv(class_names_file, header=None, names=['', 'Category'])
class_names = dict(class_names_df.set_index('').to_dict()['Category'])

# Define data transformations
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Load test dataset
test_dataset = CustomTestDataset(root_dir=test_dir, transform=transform)

# Create DataLoader for test data
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# Load pre-trained EfficientNet-B0 model
model = EfficientNet.from_pretrained('efficientnet-b1', num_classes=100)

# Load saved model parameters
model.load_state_dict(torch.load(model_path))
model.eval()

# Move model to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Initialize lists to store predictions
predictions = []
image_names = []

# Iterate over test dataset and make predictions
for images, names in test_loader:
    images = images.to(device)
    with torch.no_grad():
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        predictions.extend(predicted.cpu().numpy())
        image_names.extend(names)

# Convert predicted class numbers to class names
predicted_class_names = [class_names[class_num] for class_num in predictions]

# Create a DataFrame to store results
results_df = pd.DataFrame({'Id': [idx.split('.')[0] for idx in image_names], 
                           'Category': predicted_class_names})

# Sort the results DataFrame by the 'Id' column
results_df_sorted = results_df.sort_values(by='Id', key=lambda x: x.astype(int))

# Save results to CSV file
results_df_sorted.to_csv('test_results_orderedb1_epoch2.csv', index=False)

print("Results saved to test_results_orderedb1_epoch2.csv.")