Calculate Perceptual Similarity Loss between image in a folder

calc_percep.py

from models import dist_model as dm
from util import util
import os
import ipdb
import glob
import csv
import cv2
import skimage.io as sio

class PercepSimilarity:
    def __init__(self,
                folder_path, 
                log_file_name,
                start_idx=1, # apparently gaurav thinks indices start from 1 (moron)
                num_steps=3):
        super().__init__()
        self.model = dm.DistModel()
        self.folder_path = folder_path
        self.log_file_name = log_file_name
        if not os.path.exists(self.folder_path):
            print('No such folder')
            import ipdb; ipdb.set_trace()
        else:
            # assuming that gaurav is sane and stores images as .png
            self.filenames = glob.glob(os.path.join(self.folder_path, '*.png'))

        self.model.initialize(model='net-lin',net='alex', use_gpu=False)
        # Number of steps for generation that you have considered
        self.num_steps = num_steps
        # prefix for naming the file
        self.prefix = 'sample'
        self.start_idx = start_idx

    def calculate_loss(self, img0, img1):
        dist = self.model.forward(img0, img1)
        return dist

    def resize(self, image, size=256):
        '''need to resize, stupid gaurav
        '''
        image = cv2.resize(image, (size, size))
        return image

    def calculate(self):
        number_of_samples = len(self.filenames)
        self.num_steps = 3
        dictionary_list = []
        for i in range(self.start_idx, number_of_samples//self.num_steps):
            for j in range(self.start_idx, self.num_steps):
                dictionary = dict()
                filename1 = "_".join([self.prefix, str(i), str(j)]) + '.png'
                filename2 = "_".join([self.prefix, str(i), str(j+1)]) + '.png'
                # use the im2tensor function in the util (keep shit consistent)
                img1 =  self.resize(
                        util.load_image(os.path.join(self.folder_path, filename1)))
                img2 = self.resize(
                        util.load_image(os.path.join(self.folder_path, filename2)))
                dist = self.calculate_loss(util.im2tensor(img1),
                                           util.im2tensor(img2))
                dictionary['filename1'] = filename1
                dictionary['filename2'] = filename2
                dictionary['dist'] = dist
                dictionary_list.append(dictionary)
        with open(self.log_file_name, 'w') as csv_file:
            writer = csv.DictWriter(csv_file, dictionary.keys())
            writer.writeheader()
            for data in dictionary_list:
                writer.writerow(data)


if __name__ == '__main__':
    folder_path = '~/Downloads/output_ours'
    obj = PercepSimilarity(os.path.expanduser(folder_path),
                           './test.csv')
    obj.calculate()