Spaces:
Sleeping
Sleeping
File size: 2,734 Bytes
7fc845d |
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 95 96 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from transformers import ViTForImageClassification
from PIL import Image
import os
import pandas as pd
class CustomDataset(Dataset):
def __init__(self, dataframe, transform=None):
self.dataframe = dataframe
self.transform = transform
def __len__(self):
return len(self.dataframe)
def __getitem__(self, idx):
image_path = self.dataframe.iloc[idx, 0] # Image path is in the first column
image = Image.open(image_path).convert('RGB') # Convert to RGB format
if self.transform:
image = self.transform(image)
return image
if __name__ == "__main__":
# Check for GPU availability
device = torch.device('cuda')
# Load the pre-trained ViT model and move it to GPU
model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224').to(device)
model.classifier = nn.Linear(model.config.hidden_size, 2).to(device)
# Define the image preprocessing pipeline
preprocess = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor()
])
# Load the test dataset
### need to recive image from gratio/streamlit
test_set = 'datasets/'
image_paths = []
for filename in os.listdir(test_set):
image_paths.append(os.path.join(test_set, filename))
dataset = pd.DataFrame({'image_path': image_paths})
test_dataset = CustomDataset(dataset, transform=preprocess)
test_loader = DataLoader(test_dataset, batch_size=32)
# Load the trained model
model.load_state_dict(torch.load('trained_model.pth'))
# Evaluate the model
model.eval()
confidences = []
predicted_labels = []
with torch.no_grad():
for images in test_loader:
images = images.to(device)
outputs = model(images)
logits = outputs.logits # Extract logits from the output
probabilities = F.softmax(logits, dim=1)
confidences_per_image, predicted = torch.max(probabilities, 1)
predicted_labels.extend(predicted.cpu().numpy())
confidences.extend(confidences_per_image.cpu().numpy())
print(predicted_labels)
print(confidences)
confidence_percentages = [confidence * 100 for confidence in confidences]
for label, confidence in zip(predicted_labels, confidence_percentages):
print(f"Predicted label: {label}, Confidence: {confidence:.2f}%")
|