import streamlit as st
import torch
from PIL import Image
import torchvision.transforms as transforms
from model import SiameseNetwork 


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


model = SiameseNetwork().to(device)
model.load_state_dict(torch.load("siamese_model.pth", map_location=device))
model.eval()


transform = transforms.Compose([
    transforms.Resize((100, 100)), 
    transforms.Grayscale(num_output_channels=1),  
    transforms.ToTensor(),  # Converting image to tensor
])

# Streamlit interface
st.title("Signature Forgery Detection with Siamese Network")
st.write("Upload two signature images to check if they are from the same person or if one is forged.")

# Upload images
image1 = st.file_uploader("Upload First Signature Image", type=["png", "jpg", "jpeg"])
image2 = st.file_uploader("Upload Second Signature Image", type=["png", "jpg", "jpeg"])

if image1 and image2:
    
    img1 = Image.open(image1).convert("RGB")
    img2 = Image.open(image2).convert("RGB")

    ## Displaying input image
    col1, col2 = st.columns(2)
    with col1:
        st.image(img1, caption='First Signature Image', use_container_width=True)
    with col2:
        st.image(img2, caption='Second Signature Image', use_container_width=True)

    # Transforming the images before feeding them into the model
    img1 = transform(img1).unsqueeze(0).to(device)
    img2 = transform(img2).unsqueeze(0).to(device)

    # Predicting similarity using the Siamese model
    output1, output2 = model(img1, img2)
    euclidean_distance = torch.nn.functional.pairwise_distance(output1, output2)

    # Setting a threshold for similarity 
    threshold = 0.5 

    # Display similaritying score and interpretation
    st.success(f'Similarity Score (Euclidean Distance): {euclidean_distance.item():.4f}')
    if euclidean_distance.item() < threshold:
        st.write("The signatures are likely from the **same person**.")
    else:
        st.write("The signatures **do not match**, one might be **forged**.")