runner.py

import numpy as np
import os
import torch
from tqdm import tqdm
import zipfile

from matplotlib import pyplot as plt
import slayerpytorch as snn

from dataloader.base import mvecDatasetBase
from dataloader.base import MODDatasetBase
from dataloader.base import EVIMODatasetBase

import dataloader.base as base

from model import getNetwork

from utils.gpu import moveToGPUDevice
from utils.rbase import RBase

from torchvision import transforms
from tensorboardX import SummaryWriter
from torch.utils.data.dataloader import default_collate
import torch.nn as nn

from PIL import Image

def my_collate(batch):
    batch = list(filter (lambda x:x is not None, batch))
    if not batch:
        return None

    return default_collate(batch)

class Runner(RBase):

    def __init__(self, crop, maxBackgroundRatio, datasetType, datafile, 
                    checkpoint, modeltype, 
                    log_config, general_config, 
                    maskDir, incrementalPercent,
                    saveImages, saveImageInterval, imageDir, imageLabel=""):
        super().__init__(datafile, log_config, general_config)
       
        self.output_dir = self.log_config.getOutDir() 
        self.genconfigs = snn.params(general_config)
        self.checkpoint = checkpoint
        self.modeltype = modeltype

        self.maskDir = maskDir
        self.incrementalPercent = incrementalPercent
        self.saveImages = saveImages
        self.saveImageInterval = saveImageInterval
        self.imageDir = imageDir
        self.imageLabel = imageLabel

        if(datasetType == "EVIMO"):
            database = base.EVIMODatasetBase(datafile, self.genconfigs, self.maskDir, crop, maxBackgroundRatio, incrementalPercent)
            print("EVIMO used")
        elif(datasetType == "MOD"):
            database = base.MODDatasetBase(datafile, self.genconfigs, self.maskDir, crop, maxBackgroundRatio, incrementalPercent)
            print("MOD used")
        else:
            raise Exception("Only EVIMO or MOD datasets with hdf5 format generated by preprocessing scripts handled with this code.")
        
        
        dataset_size = len(database)
        # uncomment if you want to split test/train using single hdf5 file
        
        # test_split = self.genconfigs['model']['testSplit']
        # test_size = int(test_split * dataset_size)
        # train_size = dataset_size - test_size

        # torch.manual_seed(0)
        # train_dataset, test_dataset = torch.utils.data.random_split(database)

        num_workers = self.genconfigs['hardware']['readerThreads']
        batch_size = self.genconfigs['batchsize']

        self.loader = torch.utils.data.DataLoader(
                database,
                batch_size=batch_size,
                shuffle=False,
                num_workers=num_workers,
                pin_memory=True,
                collate_fn=my_collate,
                drop_last=False)

        self.tb_writer = SummaryWriter(self.output_dir)    

    def test(self):
        self._loadNetFromCheckpoint(self.checkpoint, self.modeltype)
        self.net = self.net.eval()

        print("saving to: ", self.output_dir)
        
        total_MSE = 0
        total_IOU = 0
        
        scalar_i = 0
        tot_frames = 0

        if self.saveImages and not os.path.exists(self.imageDir):
                os.mkdir(self.imageDir)

        with torch.no_grad():
            for i, data in enumerate(tqdm(self.loader, desc='testing')):
                if (data == None):
                    continue

                data = moveToGPUDevice(data, self.device, self.dtype)
                
                spikes_input = data['spike_tensor']
                spikes_masked = data['masked_spike_tensor']

                #model returns output and intermediate layers  
                try: 
                    spike_pred =  self.net.forward(spikes_input).to(self.device)

                #model returns just the output layer
                except: 
                    spike_pred_arr = self.net.forward(spikes_input)
                    spike_pred = spike_pred_arr[0].to(self.device)  
                    

                spike_input_crop = spikes_input[:, :, :spike_pred.size(2), :spike_pred.size(3), :spike_pred.size(4)]
                spikes_masked_crop = spikes_masked[:, :, :spike_pred.size(2), :spike_pred.size(3), :spike_pred.size(4)]

                spike_pred_2D = torch.sum(spike_pred, axis = (0,1,4))
                spike_mask_2D = torch.sum(spikes_masked_crop, axis = (0,1,4))

                #calculate metrics
                ioucriterion = snn.loss(self.genconfigs).to(self.device) 
                iou = ioucriterion.getIOU(spike_pred_2D, spike_mask_2D)
                total_IOU += iou*len(data['ratio'])

                print(i, ": ", "curr iou", iou, data['ratio'])
                self.tb_writer.add_scalar('iou', iou.item(), scalar_i)

                scalar_i += 1
                tot_frames += len(data['ratio'])
                
                if self.saveImages and (i)%self.saveImageInterval== 0:
                    spikes_inputnp = np.array(spikes_input.detach().cpu())
                    spikes_maskednp = np.array(spikes_masked.detach().cpu())
                    spikesPred_np = np.array(spike_pred.detach().cpu())

                    print("save to: ", self.output_dir)

                    for batch in range(0,spikes_input.shape[0]):
                        curr_num = data['file_number'][batch]
                        
                        im = Image.fromarray(np.uint8(np.sum(spikesPred_np[batch,:,:,:,:], axis=(0,3))*255))
                        im.save(os.path.join(self.imageDir,"_pred_epoch{}".format(curr_num) + self.imageLabel + ".jpg"))

                        im2 = Image.fromarray(np.uint8(np.sum(spikes_maskednp[batch,:,:,:,:], axis=(0,3))*255))
                        im2.save(os.path.join(self.imageDir,"_ideal_epoch{}".format(curr_num) + self.imageLabel + ".jpg"))

                        im3 = Image.fromarray(np.uint8(np.sum(spikes_inputnp[batch,:,:,:,:], axis=(0,3))*255))
                        im3.save(os.path.join(self.imageDir,"_input_epoch{}".format(curr_num)) + self.imageLabel + ".jpg")

        print("save to: ", self.output_dir)

        if self.saveImages:
            print("saving images to", os.getcwd(), self.imageDir)

        print("mean_IOU for {} batches of frames".format(tot_frames), total_IOU/tot_frames)