thotranexe/CNN

1import torch2import torch.nn as nn3import torchvision.models as models4import torchvision.transforms as transforms5from torch.autograd import Variable6from PIL import Image7from torchvision import transforms8import os9from glob import glob10import json11from json import JSONEncoder12import numpy13from sklearn.neighbors import NearestNeighbors14import streamlit as st15import dropbox16import io17 18dbx=dropbox.Dropbox("sl.BdGZL5PHc3UJpUclN48L18TtjhntJy2NiUk89HK_mylKpAw9WJH3ScGVfPMol-qapZuUjNSgaPkMlB4-rrsK9_Nx2biMeujegAmlT1GQmgA4YJNsV0AEqr--91Zt7z8es3gNTFo")19 20resnet=models.resnet50(pretrained=True)21layer = resnet._modules.get('avgpool')22#grab all images in the lfw folder23import os24from glob import glob25path="./lfw"26 27result = [y for x in os.walk(path) for y in glob(os.path.join(x[0], '*.jpg'))]28resnet.eval29 30d={}31 32preprocess=transforms.Compose([transforms.Resize(256),33                               transforms.CenterCrop(224),34                               transforms.ToTensor(),35                               transforms.Normalize(mean=[.485,.456,.406],std=[.229,.224,.225])36                               ])37 38def get_vector(image):39    # Create a PyTorch tensor with the transformed image40    t_img = preprocess(image)41    my_embedding = torch.zeros(2048)42 43    # Define a function that will copy the output of a layer44    def copy_data(m, i, o):45        my_embedding.copy_(o.flatten())                 # <-- flatten46 47    # Attach that function to our selected layer48    h = layer.register_forward_hook(copy_data)49    # Run the model on our transformed image50    with torch.no_grad():                               # <-- no_grad context51        resnet(t_img.unsqueeze(0))                       # <-- unsqueeze52    # Detach our copy function from the layer53    h.remove()54    # Return the feature vector55    return my_embedding56 57#if not d:58#    for image in result:59#        d[image]=get_vector(Image.open(image).convert('RGB')).numpy()60 61st.write("cnn assignment")62 63#class NumpyArrayEncoder(JSONEncoder):64#    def default(self, obj):65#        if isinstance(obj, numpy.ndarray):66#            return obj.tolist()67#        return JSONEncoder.default(self, obj)68 69#with open("sample.json", "w") as outfile:70#    json.dump(d, outfile,cls=NumpyArrayEncoder)71 72	73 74 75_, res = dbx.files_download("/sample.json")76 77with io.BytesIO(res.content) as stream:78    data = json.load(stream)79 80image=st.file_uploader(label="upload your own file",type="jpg")81if image is None:82    st.write("upload an image")83else:84    input=get_vector(Image.open(image).convert('RGB')).numpy()85    featurelist=[]86    for img in data:87        featurelist.append(data[img])88    neighbors = NearestNeighbors(n_neighbors=10, algorithm='brute',metric='euclidean').fit(featurelist)89    distances, indices = neighbors.kneighbors(input.reshape(1,-1))90    simular=[]91    for i in range(10):92        simular.append(result[indices[0][i]])93    st.image(simular,caption=simular)94