Analyze_func.py

#import python libraries
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import skimage as ski
import matplotlib.pyplot as plt
import numpy as np
import argparse
from sklearn.neighbors import NearestNeighbors


#import the functions from the other scripts
from HMDataLoading import HMDataset
from preprocess import *
from model import *
from train_eval import *
from plot import *

'''
# Visualize the images before and after reconstruction
Input: 
- model: trained model
- data: train/val dataset 
- idx: image to visualize
Output:
None (Show image)
'''
def Visualize_Recon(model, data, idx):
    # load dataloader
    if data == "train":
        dataset = torch.load('torch_dataset/train_dataset.pt')
        #dataloader = DataLoader(train_dataset, batch_size=1, shuffle=False)
    elif data == "val":
        dataset= torch.load('torch_dataset/val_dataset.pt')
        #dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False)
        
    # load model
    if model == "auto_paper":
        trained_model = Autoencoder_paper()
        trained_model.load_state_dict(torch.load("trained_models/autoencoder_paper_lr0.01.pt", map_location=torch.device('cpu')))
    elif model == "vae":
        trained_model = VAE()
        trained_model.load_state_dict(torch.load("trained_models/VAE_lr0.01.pt", map_location=torch.device('cpu')))
    trained_model.eval()
    
    # visualize images before and after reconstruction
    criterion = nn.MSELoss()
    with torch.no_grad():
        # format data 
        X, imgPATH = dataset[idx]
        
        X = X.float()
        if model == "vae":
            X= X.unsqueeze(0)
            org_image = X[0].numpy().transpose(1,2,0)
            X_rec, X_latent, _ , _ = trained_model(X)
            image_rec = X_rec[0].numpy().transpose(1,2,0)
        else:
            org_image = X.numpy().transpose(1,2,0)
            X_rec, X_latent = trained_model(X)
            image_rec = X_rec.numpy().transpose(1,2,0)
            
        print("showing the original image:")
        ski.io.imshow(org_image)
        ski.io.show()
                    
        print("showing the reconstructed image:")
        ski.io.imshow(image_rec)
        ski.io.show()
            
        print("MSE loss for this example:", criterion(X_rec, X))


'''
# Get the latent vector representation of the images in a dataset
Input: 
-model: model to use
-set: dataset to use
Output:
-latents: a list of latent vectors
-paths: the path of the original images
'''
def GetLatent(model, set):

    # load dataloader
    if set == "train":
        dataset = torch.load('torch_dataset/train_dataset.pt')
        dataloader = DataLoader(dataset, batch_size=len(dataset), shuffle=False)
    elif set == "val":
        dataset= torch.load('torch_dataset/val_dataset.pt')
        dataloader = DataLoader(dataset, batch_size=len(dataset), shuffle=False)
        
    # load model
    if model == "auto_paper":
        trained_model = Autoencoder_paper()
        trained_model.load_state_dict(torch.load("trained_models/autoencoder_paper_lr0.01.pt", map_location=torch.device('cpu')))
    elif model == "vae":
        trained_model = VAE()
        trained_model.load_state_dict(torch.load("trained_models/VAE_lr0.01.pt", map_location=torch.device('cpu')))
    trained_model.eval()

    # store the latent data
    with torch.no_grad():
        for data in dataloader:
            # format data 
            X, imgPATHs = data
            X = X.float()
            # model outputs
            if model == "auto_paper":
                X_rec_all, X_latent_all = trained_model(X)
            elif model == "vae":
                X_rec_all, X_latent_all, _, _ = trained_model(X)
            break
    return X_latent_all, imgPATHs



'''
# Show the image from the original data
Input:
- imgpath: the path of the preprocessed data
'''
def ShowOrgImage(imgpath):
    orgImgPath = imgpath.replace(".npy", ".jpg").replace("preprocessed_images", "images")
    
    #show image
    ski.io.imshow(orgImgPath)
    ski.io.show()





'''
# Find nearest neighbors of the images for recommendation
Input:
- model: the model to used to encode an image ["auto_paper", "vae", "original"]
- similarity: similarity matrices to measure the similarity between latent vectors
- numRank: number of products to recommend
- X_latent_all: latent vectors from the model (generated by GetLatent function)
- imgPATHs: a tuple of all the image paths in the preprocessed dataset
- idxList: list of indices of the image to recommend (corresponding to the X_latent_all and imgPATHs index)

Output:
N/A (show numRank images)
'''
def Recommend(model, similarity, numRank, X_latent_all, imgPATHs, idxList):
    
    if model == "auto_paper":
        X_latent_vec_all = torch.flatten(X_latent_all, start_dim=1).numpy()
    elif model == "vae":
        X_latent_vec_all = X_latent_all.numpy()
    
    knn = NearestNeighbors(n_neighbors=numRank+1, metric=similarity)   
    knn.fit(X_latent_vec_all) 
    distances, indices = knn.kneighbors(X_latent_vec_all)
        
    # loop over the required images 
    for requireImg_index in idxList:
        
        #show current image
        curr_imgPATH = imgPATHs[requireImg_index]
        print("current image:")
        ShowOrgImage(curr_imgPATH)
        
        #show recommended images:
        recommend_imgs_indices = indices[requireImg_index].tolist() #a list of index that are ranked by the similarity
        for i in range(numRank):
            similar_img_index = recommend_imgs_indices[i+1] #skip the 0 one since it is the same with itself
            similar_img_path = imgPATHs[similar_img_index]
            print(f"showing the rank {i+1} image")
            ShowOrgImage(similar_img_path)
            


 




if __name__ == "__main__":
    print("start")
    
    print("auto_paper:")
    X_latent_all, imgPATHs = GetLatent("auto_paper", "val")
    print(X_latent_all.size())
    print(len(imgPATHs))
    
    idxList = [3,5,7]
    Recommend(model="auto_paper", similarity="cosine", numRank=3, X_latent_all=X_latent_all, imgPATHs=imgPATHs, idxList=idxList)
    
    
    
    #print(X_latent_all[0])
    #print(imgPATHs[0])
    
    '''
    print("vae:")
    X_latent_all, imgPATHs = GetLatent("vae", "val")
    print(X_latent_all.size())
    print(len(imgPATHs))
    '''