Spaces:
Sleeping
Sleeping
| #%% | |
| import gradio as gr | |
| from PIL import Image | |
| from torchvision import transforms | |
| from siamese_nn import Siamese_nn | |
| import torch, os | |
| import torch.nn.functional as F | |
| model = Siamese_nn() | |
| weights = torch.load('trained_model', map_location=torch.device('cpu')) | |
| model.load_state_dict(weights) | |
| model.eval() | |
| file_list = os.listdir("data") | |
| examples = [] | |
| usersIndexes = [] | |
| for x in file_list: | |
| if x[0:2] not in usersIndexes: | |
| usersIndexes.append(x[0:2]) | |
| for user in usersIndexes: | |
| usersImages = [x for x in file_list if str(user) in x] | |
| notUsersImages = [x for x in file_list if str(user) not in x] | |
| for userImage in usersImages: | |
| for userImageCopy in usersImages: | |
| examples.append([userImage, userImageCopy, 0]) | |
| for notUser in notUsersImages: | |
| examples.append([userImage, notUser, 1]) | |
| def predict(input1, input2, label=None): | |
| img1_PIL = Image.open(f'data/{input1}') | |
| img2_PIL = Image.open(f'data/{input2}') | |
| img1 = transforms.ToTensor()(img1_PIL).unsqueeze(0) | |
| img2 = transforms.ToTensor()(img2_PIL).unsqueeze(0) | |
| for el in examples: | |
| if input1 == el[0] and input2 == el[1] and el[2] == 0: | |
| label = 'Scans of the same finger' | |
| break | |
| if input1 == el[0] and input2 == el[1] and el[2] == 1: | |
| label = 'Scans of different fingers' | |
| with torch.no_grad(): | |
| out1, out2 = model(img1, img2) | |
| pred = F.pairwise_distance(out1, out2) | |
| if pred < 0.6: | |
| decision = f'Access granted, confidence: {pred.item():4f}' | |
| else: | |
| decision = f'Access denied, confidence: {pred.item():4f}' | |
| return img1_PIL, img2_PIL, decision, label | |
| #%% | |
| css = """ | |
| .gradio-container { | |
| height: 100vh; | |
| max-width: 1024px !important; | |
| } | |
| .my_img { | |
| max-height: 288px !important; | |
| object-fit: cover !important; | |
| } | |
| #res div h2 { color: #07ef03; } | |
| """ | |
| js = """ | |
| () => { | |
| label = document.querySelector("#res div h2"); | |
| txt = label.textContent.split(",")[0] | |
| if (txt === 'Access granted') { | |
| label.style.color = "#07ef03"; | |
| } | |
| if (txt === 'Access denied') { | |
| label.style.color = "red"; | |
| } | |
| } | |
| """ | |
| def refresh(): | |
| image = Image.open(f'data/{file_list[0]}') | |
| return image, image | |
| with gr.Blocks(css=css, elem_classes=['container']) as demo: | |
| with gr.Row(): | |
| with gr.Row(): | |
| drop1 = gr.Dropdown(value=file_list[0], | |
| choices=file_list, | |
| label='Select first image', | |
| scale=1, | |
| ) | |
| drop2 = gr.Dropdown(value=file_list[0], | |
| choices=file_list, | |
| label='Select second image', | |
| scale=1, | |
| ) | |
| label = gr.Label(value='Scans of the same finger', show_label=False) | |
| with gr.Row(): | |
| img1 = gr.Image(height=288, # unfortunately value doesn't work properly | |
| width=256, | |
| interactive=False, | |
| scale=1, | |
| label='first image', | |
| show_download_button=False, | |
| show_share_button=False, | |
| elem_classes=['my-img']) | |
| img2 = gr.Image(height=288, | |
| width=256, | |
| interactive=False, | |
| scale=1, | |
| label='second image', | |
| show_download_button=False, | |
| show_share_button=False, | |
| elem_classes=['my-img']) | |
| output = gr.Label(value=predict(*examples[0])[2], elem_id='res', show_label=False) | |
| drop1.change(fn=predict, inputs=[drop1, drop2], outputs=[img1, img2, output, label]) | |
| drop2.change(fn=predict, inputs=[drop1, drop2], outputs=[img1, img2, output, label]) | |
| output.change(fn=None, inputs=None, js=js) | |
| # initial img load workaround | |
| demo.load(fn=refresh, inputs=None, outputs=[img1, img2]) | |
| demo.launch() | |
| # %% | |