Spaces:
Runtime error
Runtime error
File size: 2,444 Bytes
dc282ed |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
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()
|