Spaces:
Sleeping
Sleeping
| import json | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import torchvision | |
| import torchvision.transforms as transforms | |
| from torch.utils.data import DataLoader | |
| print(torch.__version__) | |
| # Define hyperparameters (adjust as needed) | |
| batch_size = 60 | |
| epochs = 10 | |
| learning_rate = 0.001 | |
| num_classes = 10 | |
| clss = [] | |
| # Define data transformations (adapt for your dataset) | |
| transform = transforms.Compose([ | |
| transforms.Resize((224, 224)), # Adjust image size if needed | |
| transforms.ToTensor(), | |
| transforms.Grayscale(), # Convert to grayscale | |
| transforms.Normalize((0.5,), (0.5,)) # Grayscale normalization (modify for RGB) | |
| ]) | |
| # Define data paths (replace with your directory structure) | |
| train_data_dir = "SOCOFing/first_ten_files" | |
| val_data_dir = "SOCOFing/first_ten_files" | |
| # Load datasets using appropriate loaders (replace with your data handling logic) | |
| train_dataset = torchvision.datasets.ImageFolder(train_data_dir, transform=transform) | |
| val_dataset = torchvision.datasets.ImageFolder(val_data_dir, transform=transform) | |
| # Create data loaders | |
| train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) | |
| val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False) | |
| # Define the PyTorch CNN model | |
| class FingerprintRecognitionModel(nn.Module): | |
| def __init__(self, num_classes): | |
| super(FingerprintRecognitionModel, self).__init__() | |
| self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1) # Grayscale input | |
| self.pool = nn.MaxPool2d(kernel_size=2, stride=2) | |
| self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) | |
| self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1) | |
| self.fc1 = nn.Linear(128 * 28 * 28, 256) # Adjust output size based on input | |
| self.fc2 = nn.Linear(256, num_classes) | |
| def forward(self, x): | |
| x = self.pool(F.relu(self.conv1(x))) | |
| x = self.pool(F.relu(self.conv2(x))) | |
| x = self.pool(F.relu(self.conv3(x))) | |
| x = x.view(-1, 128 * 28 * 28) # Flatten | |
| x = F.relu(self.fc1(x)) | |
| x = F.softmax(self.fc2(x), dim=1) | |
| return x | |
| # Create the model instance | |
| model = FingerprintRecognitionModel(num_classes) | |
| # Define loss function and optimizer (consider using more advanced techniques) | |
| criterion = nn.CrossEntropyLoss() | |
| optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) | |
| # # Train the model | |
| # for epoch in range(epochs): | |
| # for images, labels in train_loader: | |
| # # Forward pass, calculate loss | |
| # outputs = model(images) | |
| # loss = criterion(outputs, labels) | |
| # | |
| # # Backward pass and optimize | |
| # optimizer.zero_grad() | |
| # loss.backward() | |
| # optimizer.step() | |
| # | |
| # # (Optional) Print training progress (replace with your logging) | |
| # if (epoch + 1) % 100 == 0: # Print every 100 mini-batches | |
| # print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}') | |
| # Train the model | |
| for epoch in range(epochs): | |
| for images, labels in train_loader: | |
| # Forward pass, calculate loss | |
| outputs = model(images) | |
| loss = criterion(outputs, labels) | |
| # Backward pass and optimize | |
| optimizer.zero_grad() | |
| loss.backward() | |
| optimizer.step() | |
| # (Optional) Print training progress (replace with your logging) | |
| if (epoch + 1) % 1 == 0: # Print every 100 mini-batches | |
| print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.4f}') | |
| # Access class indices from the training data generator | |
| class_indices = train_loader.dataset.class_to_idx | |
| # Invert the dictionary to get labels from indices (optional) | |
| labels = dict((v, k) for k, v in class_indices.items()) | |
| print(f'Class Labels: {labels}') | |
| clss.append(labels) | |
| # Save class labels to a JSON file (example) | |
| with open('class_labels.json', 'w') as f: | |
| json.dump(labels, f) | |
| # Test the model (replace with your evaluation metrics) | |
| with torch.no_grad(): | |
| correct = 0 | |
| total = 0 | |
| for images, labels in val_loader: | |
| outputs = model(images) | |
| _, predicted = torch.max(outputs.data, 1) | |
| total += labels.size(0) | |
| correct += (predicted == labels).sum().item() | |
| print(f'Test Accuracy: {100 * correct / total:.2f}%') | |
| # (Optional) Save the model (consider using serialization libraries like TorchScript) | |
| torch.save(model.state_dict(), 'fingerprint_recognition_model.pt') | |
| print(clss) | |