Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import streamlit as st | |
| import torch.nn.functional as F | |
| import torch | |
| import torch.nn as nn | |
| import torchvision.transforms as transforms | |
| from PIL import Image | |
| import cv2 | |
| num_classes = 10 | |
| # Class definition for the model (replace with your actual model definition) | |
| 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 | |
| # Load the model (replace with your path) | |
| model_path = 'fingerprint_recognition_model.pt' | |
| model = FingerprintRecognitionModel(num_classes) | |
| model.load_state_dict(torch.load(model_path)) | |
| model.eval() | |
| # Function to preprocess the image (assuming grayscale) | |
| # def preprocess_image(image): | |
| # img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) # Convert to grayscale | |
| # img_resized = cv2.resize(img, (224, 224)) | |
| # transform = transforms.Compose([ | |
| # transforms.ToTensor(), | |
| # transforms.Normalize((0.5,), (0.5,)) | |
| # ]) | |
| # img_tensor = transform(img_resized).unsqueeze(0) # Add a batch dimension | |
| # return img_tensor | |
| def preprocess_image(image): | |
| """Preprocesses an image for classification with a PyTorch model. | |
| Args: | |
| image: The image to preprocess. This can be a NumPy array representing | |
| a color or grayscale image. | |
| Returns: | |
| A PyTorch tensor representing the preprocessed image with a batch dimension. | |
| """ | |
| # Ensure image is a NumPy array | |
| if not isinstance(image, np.ndarray): | |
| image = np.array(image) # Convert if necessary | |
| # Convert to grayscale if color image | |
| if len(image.shape) == 3 and image.shape[2] == 3: # Check for color image (3 channels) | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) | |
| dsize = (224, 224) # Explicitly define output size | |
| img_resized = cv2.resize(image, dsize) | |
| # Create a PyTorch transform (assuming model expects normalized input) | |
| transform = transforms.Compose([ | |
| transforms.ToTensor(), # Convert to PyTorch tensor | |
| transforms.Normalize((0.5,), (0.5,)) # Normalize using mean and standard deviation | |
| ]) | |
| # Apply transforms and add batch dimension | |
| img_tensor = transform(img_resized).unsqueeze(0) | |
| return img_tensor | |
| # Predict class for a new image | |
| def predict_class(image): | |
| img_tensor = preprocess_image(image) | |
| with torch.no_grad(): | |
| outputs = model(img_tensor) | |
| _, predicted = torch.max(outputs.data, 1) | |
| predicted_class = int(predicted.item()) | |
| return predicted_class | |
| def main(): | |
| """ Streamlit App for Fingerprint Classification""" | |
| st.title("Fingerprint Recognition - Class Prediction") | |
| uploaded_file = st.file_uploader("Choose a Fingerprint Image", type=['png', 'jpg', 'jpeg','bmp']) | |
| if uploaded_file is not None: | |
| # Convert uploaded image to cv2 format | |
| image = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), cv2.IMREAD_COLOR) | |
| st.image(image, caption='Uploaded Fingerprint', use_column_width=True) | |
| # Preprocess and predict class | |
| processed_img = preprocess_image(image) | |
| predicted_class = predict_class(processed_img) | |
| # Display | |
| st.write(f"Predicted Class: {predicted_class}") | |
| class_labels = {0: "Left_index_finger", 1: "Left_little_finger", 2: "Left_middle_finger", 3: "Left_ring_finger", | |
| 4: "Left_thumb_finger", 5: "Right_index_finger", 6: "Right_little_finger", 7: "Right_middle_finger", | |
| 8: "Right_ring_finger", 9: "Right_thumb_finger"} | |
| st.write(f"Class Label: {class_labels[predicted_class]}") | |
| main() | |