import io import tarfile import numpy as np import random as RAND import torch import torchvision.transforms as TRNSFM import torchvision.models as MDLS from PIL import Image as IMG from scipy.spatial.distance import cosine import streamlit as st def similar(image): # Function for Streamlit App pict = form(IMG.open(image).convert('RGB')) pictFeats = mod(pict.unsqueeze(0)).detach().numpy().squeeze() for na, feat in feats.items(): s = 1 - cosine(pictFeats, feat) simScores.append((na, s)) simScores.sort(key=lambda x: x[1], reverse=True) st.write("### Selected Image") test = IMG.open(image) test.show() print('\n') # 10 Most Similar Images from Dictionary st.write("### 10 Most Similar Images") for na in simScores[:10]: for x in range(10): st.write(f"### {x}") new_na = na[:3] + "images/" + na[3:] new_path = "http://vis-www.cs.umass.edu/" + new_na simImages = IMG.open(new_path) simImages.show() mod = MDLS.resnet50(pretrained=True) mod.eval() mod = torch.nn.Sequential(*list(mod.children())[:-1]) inFile = tarfile.open('./datasets/lfw.tar', 'r') feats = {} simScores = [] # Similarity Scores for Later form = TRNSFM.Compose([ TRNSFM.Resize(256), TRNSFM.CenterCrop(224), TRNSFM.ToTensor(), TRNSFM.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]) ]) # Image Pre-processing stuffs = inFile.getmembers() for stuff in stuffs: # Going through the TAR file f = inFile.extractfile(stuff) if stuff.isdir(): continue if stuff.name.lower().endswith(('.jpg', '.jpeg', '.png')): n = stuff.name pic = form(IMG.open(io.BytesIO(f.read())).convert('RGB')) # Pre-processes the image before feeding it into the model feats[n] = mod(pic.unsqueeze(0)).detach().numpy().squeeze() # Stuff for App st.title("Similar Image Finder") upload = st.file_uploader("Upload an Image...", type=['.jpg', '.jpeg', '.png']) if upload is not None: similar(upload) st.write("## OR") # Random Image Selector from 5 Pictures randImages = [ './datasets/random-images/img1.jpg', './datasets/random-images/img2.jpg', './datasets/random-images/img3.jpg', './datasets/random-images/img4.jpg', './datasets/random-images/img5.jpg' ] if st.button("Surprise Me!"): # Button imageOptOne = RAND.choice(randImages) similar(imageOptOne) inFile.close()