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()