<file_sep>import numpy as np
import os
import os.path as osp
import random
import sys
import torch
import yaml
from utils.dotdict import DotDict
from data.ImageDataset import ImageDataset
import utils.im_manipulation as ImageManipulator
import matplotlib.pyplot as plt
from torchvision.utils import save_image, make_grid
from utils.Phase import Phase
import utils as utils
from torch.utils.data import DataLoader
from model.trainer import HierarchicalSRTrainer
from collections import defaultdict
from time import  time

configFile="../config/prosrgan.yaml" #Config File
n_cpu=8 #number of cpu threads to use during batch generation
file_separator = "/"


def modelling(args):

    training_dataset = ImageDataset(phase=Phase.TRAIN,
                                    img_dir= args.train.dataset.path,
                                    img_type=args.xtra.img_type,
                                    mean = args.train.dataset.mean,
                                    stddev = args.train.dataset.stddev,
                                    scales = args.data.scale,
                                    high_res=args.xtra.high_res)

    dataloader = DataLoader(
        training_dataset,
        batch_size=args.train.batch_size,
        shuffle=True,
        num_workers=n_cpu,
    )

    test_dataset = ImageDataset(phase=Phase.TEST,
                                    img_dir=args.test.dataset.path,
                                    img_type=args.xtra.img_type,
                                    mean=args.test.dataset.mean,
                                    stddev=args.train.dataset.stddev,
                                    scales=args.data.scale,
                                    high_res=args.xtra.high_res)

    testing_data_loader = DataLoader(
        test_dataset,
        batch_size=1,
        shuffle=True,
    )

    # checkpoint is the directory where checkpoints are read from and stored
    trainer = HierarchicalSRTrainer(args, dataloader, training_dataset, save_dir=args.xtra.out_path)

    best_saved = 0

    steps_per_epoch = len(training_dataset)
    errors_accum = defaultdict(list)
    errors_accum_prev = defaultdict(lambda: 0)
    total_steps = 0

    trainer.testcuda()

    print("RIKI: BEGIN EPOCHS....")

    epoch = 0

    #for epoch in range(0, args.train.epochs):
    while epoch < args.train.epochs:
        iter_start_time = time()
        trainer.set_train()

        intermediate_steps = 0
        recalibrating = False

        for i, imgs in enumerate(dataloader):
            '''if (i+1)% freq == 0:
                print("RIKI: EPOCH:%d BATCH: %d SCALE: %d BLEND: %f" % ( epoch, i, int(imgs['scale'][0].item()), trainer.blend))'''

            trainer.set_input(imgs['lr'], imgs['hr'], imgs['bicubic'], imgs['scale'][0].item())

            # GENERATOR MAKING HR IMAGE
            gen_hr = trainer.forward()

            # SCALE SWITCHING HAPPENS HERE IF #EPOCHS IS REACHED
            # L1 LOSS ALSO CALCULATED HERE - INSIDE backward()
            ret = trainer.optimize_parameters()

            if ret:
                print("SWITCHING GEARS...RESTARTING EPOCH "+str(epoch))
                total_steps-=intermediate_steps
                intermediate_steps = 0
                epoch-=1
                recalibrating = True
                break

            #dict: eg. 'l1_x2'
            errors = trainer.get_current_errors()
            for key, item in errors.items():
                errors_accum[key].append(item)

            total_steps += 1
            intermediate_steps += 1


            # HOW OFTEN TO PRINT ERROR STATUS TO CONSOLE: print_errors_freq
            if total_steps % args.train.io.print_errors_freq == 0:
                for key, item in errors.items():
                    if len(errors_accum[key]):
                        errors_accum[key] = np.nanmean(errors_accum[key])
                    if np.isnan(errors_accum[key]):
                        errors_accum[key] = errors_accum_prev[key]
                message = 'RIKI: EPOCH:%d BATCH: %d SCALE: %d BLEND: %f ERRORS: %s' % (epoch, i, int(imgs['scale'][0].item()), trainer.blend, str(errors_accum[key]))
                print(message)
                errors_accum_prev = errors_accum
                #t = time() - iter_start_time
                #iter_start_time = time()
                errors_accum = defaultdict(list)

            if total_steps % args.train.io.save_img_freq == 0:
                key='l1_x'+str(trainer.model_scale)

                # SAVING IMAGES IN BETWEEN EPOCHS
                op = ImageManipulator.combine_images(imgs['lr'], imgs['hr'], gen_hr, trainer.model_scale)
                fname = args.xtra.save_path+file_separator+"opimg_x"+str(trainer.model_scale)+"_"+str(epoch)+"_"+str(total_steps)+".jpg"
                print('RIKI: SAVE TO ' + fname)
                save_image(op, fname, normalize=True)

        epoch += 1
        if recalibrating:
            continue

        ################# update learning rate  #################
        # IF THE BEST EPOCH WAS SEEN A WHILE BACK...IT'S TIME TO UPDATE THE LEARNING RATE BY 1/2
        if (epoch - trainer.best_epoch) > args.train.lr_schedule_patience:
            #trainer.save('last_lr_%g' % trainer.lr, epoch, trainer.lr)
            print("RIKI: UPDATING LR...")
            trainer.update_learning_rate()

        #test_dataset.set_scales(trainer.current_scale_id)

        print("RIKI: BEGIN EVALUATION... "+str(trainer.current_scale_id)+" "+str(training_dataset.current_scale_id)+" "+str(training_dataset.current_scale))
        if len(testing_data_loader)>0:
            with torch.no_grad():

                # use validation set
                trainer.set_eval()
                trainer.reset_eval_result()

                test_dataset.copy_scales(training_dataset.current_scale)
                #print("RIKI: CHECK "+str(test_dataset.current_scale))
                for i, imgs in enumerate(testing_data_loader):
                    trainer.set_input(imgs['lr'],imgs['hr'],imgs['bicubic'],imgs['scale'])
                    #PSNR COMPUTATIONS HERE
                    trainer.evaluate()

                test_result = trainer.get_current_eval_result()

                trainer.update_best_eval_result(epoch, test_result)
                print(
                    'RIKI: eval at epoch %d : ' % epoch + ' | '.join([
                        '{}: {:.07f}'.format(k, v)
                        for k, v in test_result.items()
                    ]))

                print(
                    'RIKI: best so far %d : ' % trainer.best_epoch + ' | '.join([
                        '{}: {:.07f}'.format(k, v)
                        for k, v in trainer.best_eval.items()
                    ]))

                # IF CURRENT EPOCH IS THE BEST EPOCH SO FAR
                if trainer.best_epoch == epoch:
                    if len(trainer.best_eval) > 1:

                        # select only upto current training scale
                        best_key = ["psnr_x%d" % trainer.opt.data.scale[s_idx]
                                    for s_idx in range(trainer.current_scale_id + 1)]
                        best_key = [k for k in best_key
                                    if trainer.best_eval[k] == test_result[k]]

                    else:
                        best_key = list(trainer.best_eval.keys())
                    print('RIKI: BEST KEY: '+str(epoch) + '__'+ str(best_key))

                    # COPY trainer to best_trainer
                    trainer.best_trainer = trainer.save(epoch, trainer.lr, trainer.current_scale_id, False)

        if epoch % args.train.io.save_model_freq == 0:
            print('RIKI: STORING THE MODEL....AFTER A WHILE')
            trainer.store()





if __name__ == '__main__':

    with open(configFile) as file:
        try:
            params = yaml.load(file, Loader=yaml.FullLoader)
        except yaml.YAMLError as exc:
            print(exc)
            sys.exit(0)

    params = DotDict(params)

    if not osp.isdir(params.xtra.out_path):
        os.makedirs(params.xtra.out_path)

    #SAVING PARAMETERS FOR RESTART....NEEDS WORK
    #np.save(osp.join(params.ip.out_path, 'params'), params)

    experiment_id = osp.basename(params.xtra.out_path)

    print('experiment ID: {}'.format(experiment_id))

    #pprint(params)
    modelling(params)<file_sep>import numpy as np
import os
import os.path as osp
import random
import sys
import torch
import yaml
from utils.dotdict import DotDict
from data.ImageDataset import ImageDataset
import utils.im_manipulation as ImageManipulator
import matplotlib.pyplot as plt
from torchvision.utils import save_image, make_grid
from utils.Phase import Phase
import utils as utils
from torch.utils.data import DataLoader

configFile="../config/prosrgan.yaml" #Config File
n_cpu=8 #number of cpu threads to use during batch generation


def modelling(args):
    seed = args.xtra.seed
    print(args.xtra.out_path)
    print(args.test.dataset.path)
    #multiplier between lr and hr
    current_scale_idx=2



def meth(num_init_features, bn_size,
                 growth_rate, ps_woReLU, level_config, level_compression,
                 res_factor, max_num_feature, max_scale, **kwargs):
    print("HERE")



if __name__ == '__main__':

    with open(configFile) as file:
        try:
            params = yaml.load(file, Loader=yaml.FullLoader)
        except yaml.YAMLError as exc:
            print(exc)
            sys.exit(0)

    params = DotDict(params)

    if not osp.isdir(params.xtra.out_path):
        os.makedirs(params.xtra.out_path)

    #SAVING PARAMETERS FOR RESTART....NEEDS WORK
    #np.save(osp.join(params.ip.out_path, 'params'), params)

    experiment_id = osp.basename(params.xtra.out_path)

    print('experiment ID: {}'.format(experiment_id))

    '''if args.visdom:
        from lib.prosr.visualizer import Visualizer
        visualizer = Visualizer(experiment_id, port=args.visdom_port)'''

    #pprint(params)
    modelling(params)
    params.G.max_scale = max(params.data.scale)
    meth(**params.G)<file_sep>from enum import Enum

class Phase(Enum):
    TRAIN = 'train'
    VAL = 'val'
    TEST = 'test'
<file_sep>
import os
import torch
import numpy as np
from collections import OrderedDict
from utils.im_manipulation import tensor2im,eval_psnr_and_ssim
from model.gen_dis import HierarchicalGenerator
from data import ImageDataset
import utils.im_manipulation as ImageManipulator
from torchvision.utils import save_image, make_grid

class HierarchicalSRTrainer(object):
    def __init__(self,opt, data_loader, dataset:ImageDataset, save_dir="data/checkpoints"):
        super(HierarchicalSRTrainer, self).__init__()


        self.input = torch.zeros(
            opt.train.batch_size, 3, 48, 48,
            dtype=torch.float32)
        self.label = torch.zeros_like(
            self.input, dtype=torch.float32)
        self.interpolated = torch.zeros_like(
            self.label, dtype=torch.float32)

        self.opt = opt
        self.save_dir = save_dir

        self.dataset = dataset
        # NUMBER OF BATCHES POSSIBLE IN 1 RUN OF DATASET
        self.dataLen = len(data_loader)
        self.start_epoch = 0
        self.progress = self.start_epoch / opt.train.epochs
        self.blend = 1
        # INDEX OF THE CURRENT MODEL SCALE
        self.current_scale_id=0
        self.model_scale = self.opt.data.scale[self.current_scale_id]

        # Dictionary of scalewise best evaluated scores
        self.best_eval = OrderedDict([('psnr_x%d' % s, 0.0) for s in opt.data.scale])
        # Dictionary of scalewise all evaluated scores
        self.eval_dict = OrderedDict([('psnr_x%d' % s, []) for s in opt.data.scale])

        # TENSOR TO NUMPY ARRAY
        self.tensor2im = lambda t: tensor2im(t, mean=opt.train.dataset.mean,stddev=opt.train.dataset.stddev)

        opt.G.max_scale = max(opt.data.scale)

        # INITIALIZING THE GENERATOR
        self.net_G = HierarchicalGenerator(**opt.G).cuda()
        self.best_epoch = 0

        self.optimizer_G = torch.optim.Adam(
            [p for p in self.net_G.parameters() if p.requires_grad],
            lr=self.opt.train.lr,
            betas=(0.9, 0.999),
            eps=1.0e-08)
        self.lr = self.opt.train.lr

        self.l1_criterion = torch.nn.L1Loss()

        self.best_trainer = self.save(0, 0, 0, False)

        # INITIALIZATION OF SCALE & NETWORK PARAMETERS
        #self.reset_curriculum_for_dataloader()

    def testcuda(self):
        cuda = torch.cuda.is_available()
        print("GPU AVAILABLE:" + str(cuda))
        self.input = self.input.cuda(non_blocking=True)
        self.label = self.label.cuda(non_blocking=True)
        self.interpolated = self.interpolated.cuda(non_blocking=True)
        self.net_G = self.net_G.cuda()
        self.l1_criterion = self.l1_criterion.cuda()

    # THE LR AND THE BICUBIC INTERPOLATED HR IMAGE AS INPUT
    # RETURNS THE GENERATED HR OUTPUT
    def forward(self):
        # GETTING THE CURRENT BLEND VALUE
        if self.current_scale_id != 0:
            lo, hi = self.opt.train.growing_steps[self.current_scale_id * 2 - 2 : self.current_scale_id * 2]
            self.blend = min((self.progress - lo) / (hi - lo), 1)
            assert self.blend >= 0 and self.blend <= 1
        else:
            self.blend = 1

        self.output = self.net_G(self.input, upscale_factor=self.model_scale, blend = self.blend) + self.interpolated
        return self.output

    def set_input(self, lr, hr, bic, scale):
        self.input.resize_(lr.size()).copy_(lr)
        self.label.resize_(hr.size()).copy_(hr)
        self.interpolated.resize_(bic.size()).copy_(bic)
        #self.model_scale = scale

    def set_train(self):
        self.net_G.train()
        self.isTrain = True

    def backward(self):
        self.compute_loss()
        self.loss.backward()

    # GENERATOR LOSS
    def compute_loss(self):
        self.loss = 0
        self.l1_loss = self.l1_criterion(self.output, self.label)
        self.loss += self.l1_loss

    def optimize_parameters(self):
        self.optimizer_G.zero_grad()
        self.backward()
        self.optimizer_G.step()
        # update progress
        ret = self.increment_training_progress()
        return ret

    # THIS HELPS JUMP THE SCALE DURING TRAINING
    def increment_training_progress(self):
        ret = False
        """increment self.progress and D, G scale_idx"""
        # 1/2/3 = 1/(2*3) = .166
        # PITSTOPS AT 12%, 25%, 45%, 60%, 100%
        self.progress += 1 / self.dataLen / self.opt.train.epochs

        # IF THE #EPOCHS HAS HIT A BREAK-POINT - THIS MEANS A SCALE CHANGE IS TO BE INITIATED
        if self.progress > self.opt.train.growing_steps[self.current_scale_id * 2]:
            if self.current_scale_id < len(self.opt.data.scale) - 1:
                # SAVING CURRENT BEST
                print('RIKI: STORING CURRENT BEST MODEL....DUE TO SCALE UPDATE ')
                self.store()


                print('RIKI: trainer TIME TO INCREASE TRAINING SCALE')
                self.current_scale_id += 1
                self.net_G.current_scale_idx = self.current_scale_id
                self.model_scale*=2

                # SO THAT THE DATA LOADER LOADS THE CORRECT IMAGES
                self.dataset.set_scales(self.current_scale_id)
                # REPEAT self.opt.data.scale[i] FOR self.current_scale_idx + 1 TIMES
                training_scales = [
                    self.opt.data.scale[i]
                    for i in range(self.current_scale_id + 1)
                ]

                print('RIKI: UPDATED TRAINING SCALES: {}'.format(str(training_scales)))
                ret = True
            elif self.current_scale_id == len(self.opt.data.scale) - 1:
                print("FINISHED ALL EPOCHS AND SCALES...")

        return ret


    def get_current_errors(self):
        d = OrderedDict()
        if hasattr(self, 'l1_loss'):
            d['l1_x%d' % self.model_scale] = self.l1_loss.item()

        return d

    def set_learning_rate(self, lr, optimizer):
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

    def update_learning_rate(self):
        """update learning rate with exponential decay"""
        lr = self.lr * self.opt.train.lr_decay
        if lr < self.opt.train.smallest_lr:
            return
        self.set_learning_rate(lr, self.optimizer_G)
        print('update learning rate: %f -> %f' % (self.lr, lr))
        self.lr = lr

    # MODEL EVALUATION
    def evaluate(self):

        self.output = self.net_G(self.input, upscale_factor=self.model_scale,
                                 blend=self.blend) + self.interpolated

        im1 = self.tensor2im(self.label)
        im2 = self.tensor2im(self.output)
        eval_res = {
            'psnr_x%d' % self.model_scale:
                eval_psnr_and_ssim(im1, im2, self.model_scale)[0]
        }
        for k, v in eval_res.items():
            self.eval_dict[k].append(v)
        return eval_res

    # TEST EVALUATION AND IMAGE GENERATION
    def evaluate_and_generate(self, act_fname, output_image_path):

        self.output = self.net_G(self.input, upscale_factor=self.model_scale,
                                 blend=self.blend) + self.interpolated

        im1 = self.tensor2im(self.label)
        im2 = self.tensor2im(self.output)

        err = eval_psnr_and_ssim(im1, im2, self.model_scale)[0]
        eval_res = {
            'psnr_x%d' % self.model_scale: err
        }

        err_str = '{:.02f}'.format(err)
        op = ImageManipulator.combine_images_test(self.input, self.label, self.interpolated, self.output, self.model_scale)
        fname = output_image_path + self.opt.xtra.file_sep + act_fname+"_x" + str(self.model_scale) + "_"+err_str+ ".jpg"
        print('RIKI: SAVE TO ' + fname)
        save_image(op, fname, normalize=True)

        return eval_res

    def set_eval(self):
        self.net_G.eval()
        self.isTrain = False

    def reset_eval_result(self):
        for k in self.eval_dict:
            self.eval_dict[k].clear()

    # AVERAGING CURRENT PSNRs FROM eval_dict
    def get_current_eval_result(self):
        eval_result = OrderedDict()
        for k, vs in self.eval_dict.items():
            eval_result[k] = 0
            if vs:
                for v in vs:
                    eval_result[k] += v
                eval_result[k] /= len(vs)
        return eval_result

    def update_best_eval_result(self, epoch, current_eval_result=None):
        if current_eval_result is None:
            eval_result = self.get_current_eval_result()
        else:
            eval_result = current_eval_result
        is_best_sofar = any(
            [np.round(eval_result[k],2) > np.round(v,2) for k, v in self.best_eval.items()])
        #print("RIKI: trainer IS BEST SO FAR: "+str(is_best_sofar))
        if is_best_sofar:
            self.best_epoch = epoch
            self.best_eval = {
                k: max(self.best_eval[k], eval_result[k])
                for k in self.best_eval
            }

    # SAVING TO DISK
    def store(self):

        savepath = self.best_trainer['network']['path']
        torch.save(self.best_trainer['network'], savepath)

        savepath1 = self.best_trainer['optim']['path']
        torch.save(self.best_trainer['optim'], savepath1)

        print("RIKI: STORED "+savepath+" "+savepath1)

    # SAVING TO MEMORY
    def save(self, epoch, lr, scale, make_save=False):

        to_save = {
            'network': self.save_network(self.net_G, 'G', str(epoch), str(scale), make_save),
            'optim': self.save_optimizer(self.optimizer_G, 'G', epoch, lr, str(scale), make_save),
        }

        print("RIKI: SAVED LATEST MODEL %d %d %d...DISK: %s"%(epoch, lr, scale, str(make_save)))
        return to_save

    def save_network(self, network, network_label, epoch_label, scale, make_save=False):
        network = network.module if isinstance(
            network, torch.nn.DataParallel) else network
        #save_filename = '%s_net_%s_x%s.pth' % (epoch_label, network_label, scale)
        save_filename = 'net_%s_x%s.pth' % (network_label, scale)
        save_path = os.path.join(self.save_dir, save_filename)
        to_save = {
            'state_dict': network.state_dict(),
            'path': save_path
        }

        if make_save:
            torch.save(to_save, save_path)

        return to_save

    def save_optimizer(self, optimizer, network_label, epoch, lr, scale, make_save=False):
        save_filename = 'optim_%s_x%s.pth' % (network_label, scale)
        save_path = os.path.join(self.save_dir, save_filename)

        to_save = {
            'state_dict': optimizer.state_dict(),
            'epoch': epoch,
            'lr': lr,
            'path': save_path
        }

        if make_save:
            torch.save(to_save, save_path)

        return to_save

    def load(self, resume_from):
        self.load_network(self.net_G, resume_from[0])
        self.load_optimizer(self.optimizer_G, resume_from[1])

    def load_network(self, network, saved_path):
        network = network.module if isinstance(
            network, torch.nn.DataParallel) else network
        loaded_state = torch.load(saved_path)['state_dict']
        loaded_param_names = set(loaded_state.keys())

        # allow loaded states to contain keys that don't exist in current model
        # by trimming these keys;
        own_state = network.state_dict()
        extra = loaded_param_names - set(own_state.keys())
        if len(extra) > 0:
            print('Dropping ' + str(extra) + ' from loaded states')
        for k in extra:
            del loaded_state[k]

        try:
            network.load_state_dict(loaded_state)
        except KeyError as e:
            print(e)
        print('RIKI: loaded network state from ' + saved_path)

    def load_optimizer(self, optimizer, saved_path):

        data = torch.load(saved_path)
        loaded_state = data['state_dict']
        optimizer.load_state_dict(loaded_state)

        # Load more params
        self.start_epoch = data['epoch']
        self.lr = data['lr']

        print('RIKI: loaded optimizer state from ' + saved_path)<file_sep>#LOADING THE IMAGE FILES IN A CUSTOM WAY
from __future__ import print_function, division
import os
import torch
import pandas as pd
from skimage import io, transform
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
from collections.abc import Iterable
from utils.Phase import Phase

from PIL import Image
import utils.im_manipulation as ImageManipulator
#IMAGES WILL BE pixel_res
pixel_res = 32

class ImageDataset(Dataset):

    # READS ALL IMAGE FILE IN THE TARGET DIRECTORY AND CREATES TWO LISTS - KEYS & FILEPATHS WHOSE INDICES CORRESPOND
    def constructImageDictionary(self, img_dir, ext):
        img_keys = []
        img_paths = []
        for roots, dir, files in os.walk(img_dir):
            for file in files:
                if ext in file:
                    file_abs_path = os.path.join(roots, file)
                    tokens = file_abs_path.split('/');
                    ln = len(tokens)
                    geohash = tokens[ln - 2]

                    f_tokens = file.split('$')
                    date = f_tokens[0]

                    img_keys.append(geohash + '$' + date)
                    img_paths.append(file_abs_path)
        return img_keys, img_paths

    # IMAGE UPSCALING
    # img: PIL image
    # INCREASE RESOLUTION OF IMAGE BY FACTOR OF 2^pow
    def upscaleImage(self, img, pow, method=Image.BICUBIC):
        ow, oh = img.size
        scale = 2**pow
        h = int(round(oh * scale))
        w = int(round(ow * scale))
        return img.resize((w, h), method)




    def imageResize(self, image):
        new_image = image.resize((pixel_res, pixel_res))
        return new_image


    # RETURNS A LOW RES AND HIGH RES VERSION OF THE IMAGE AT A GIVEN INDEX
    def __getitem__(self, index):
        ret_data = {}
        ret_data['scale'] = self.current_scale

        scale = self.current_scale

        # Load HIGH-RES image
        if len(self.image_fnames):
            # FULL HIGH RESOLUTION IMAGE
            target_img = ImageManipulator.image_loader(self.image_fnames[index])

            w, h = target_img.size
            # IF IMAGE PIXELS IS NOT DIVISIBLE BY SCALE, CROP THE BALANCE
            target_img = target_img.crop((0, 0, w - w % scale, h - h % scale))

            target_img = ImageManipulator.imageResize(target_img, self.highest_res)

            # THIS IS HOW MUCH TO DE-MAGNIFY THE MAIN IMAGE
            # eg. if target is x2, target has to be divided by 4 and source has to be divided by 8
            div_factor = self.max_scale / self.current_scale

            # LOW-RES IMAGE GENERATION

            hr_image = ImageManipulator.downscaleImage(target_img, div_factor)
            lr_image = ImageManipulator.downscaleImage(target_img, self.max_scale)

            if self.augmentON:
                # ROTATING & FLIPPING IMAGE PAIRS
                ret_data['lr'], ret_data['hr'] = ImageManipulator.augment_pairs(
                                                            lr_image, hr_image)

            else:
                ret_data['hr'] = hr_image
                ret_data['lr'] = lr_image

            # IMAGE & FILE_NAME
            ret_data['bicubic'] = ImageManipulator.downscaleImage(ret_data['lr'],1 / scale)
            ret_data['hr_fname'] = self.image_fnames[index]

            ret_data['hr'] = self.normalize_fn(ret_data['hr'])
            ret_data['lr'] = self.normalize_fn(ret_data['lr'])
            ret_data['bicubic'] = self.normalize_fn(ret_data['bicubic'])


        return ret_data

    def __len__(self):
        return len(self.image_fnames)

    # DISPLAY AN IMAGE AT A GIVEN PATH
    def displayImage(self, image_path):
        img = Image.open(image_path)
        plt.figure()
        plt.imshow(img)
        plt.show()

    # DISPLAY AN IMAGE OBJECT
    def displayImageMem(self, img):
        plt.figure()
        plt.imshow(img)
        plt.show()

    def getTransforms(self):
        transform_list = []

        transform_list += [transforms.ToTensor(),
                           transforms.Normalize((0.5, 0.5, 0.5),
                                                (0.5, 0.5, 0.5))]
        return transforms.Compose(transform_list)

    # CALLED DURING SWITCHING OF GEARS
    def set_scales(self, s_ind):
        if s_ind >= 0 and s_ind< len(self.scales):
            self.current_scale_id = s_ind
            self.current_scale = self.scales[s_ind]
            print("RESET TO "+str(self.current_scale))
        else:
            print("RIKI: SCALE ERROR")

    def copy_scales(self, scal):
        if self.scales.index(scal) >= 0 :
            self.current_scale = scal
            #print("RESET TEST SCALE TO "+str(self.current_scale))
        else:
            print("RIKI: SCALE ERROR")

    def __init__(self, phase, img_dir, img_type, mean, stddev, scales, high_res):
        self.phase = phase
        self.mean = mean
        self.stddev = stddev

        # THE RESOLUTION OF THE x8 ACTUAL IMAGE
        self.highest_res = high_res

        # INITIALIZING SCALES
        self.scales = scales if isinstance(scales, Iterable) else [scales]
        self.current_scale = self.scales[0]
        self.current_scale_id = self.scales.index(self.current_scale)+1
        self.max_scale = np.max(self.scales)

        # IF PHASE IS TRAINING PHASE, WE AUGMENT THE INPUT IMAGE RANDOMLY
        #self.augmentON = self.phase == Phase.TRAIN
        self.augmentON = False
        # THE DIRECTIRY WHERE IMAGE ARE
        self.img_dir = img_dir

        # ALL IMAGE FILES IN AN ARRAY
        self.image_fnames = ImageManipulator.get_filenames(img_dir, img_type)

        # Input normalization
        self.normalize_fn = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(self.mean, self.stddev)
        ])
<file_sep>import torch
import torch.optim as opt
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
import numpy as np
from torchvision.models import vgg19
from math import log2
from model.layers import PixelShuffleUpsampler,Conv2d,CompressionBlock,init_weights,DenseConnectedBlock
from collections import OrderedDict

class HierarchicalGenerator(nn.Module):

    def __init__(self, num_init_features, bn_size, growth_rate, ps_woReLU, level_config,
                 res_factor, max_num_feature, max_scale, **kwargs):
        super(HierarchicalGenerator,self).__init__()
        self.max_upscale = max_scale
        self.num_pyramids = int(log2(self.max_upscale))

        self.current_scale_idx = self.num_pyramids - 1

        self.Upsampler = PixelShuffleUpsampler
        self.upsample_args = {'woReLU': ps_woReLU}

        denseblock_params = {
            'num_layers': None,
            'num_input_features': num_init_features,
            'bn_size': bn_size,
            'growth_rate': growth_rate,
        }

        num_features = denseblock_params['num_input_features']  # 160

        # Initiate network

        # each scale has its own init_conv - V ******************************************************
        for s in range(1, self.num_pyramids + 1):
            # in_channel=3, out_channel=160
            self.add_module('init_conv_%d' % s, Conv2d(3, num_init_features, 3))

        # Each denseblock forms a pyramid - 0,1,2
        for i in range(self.num_pyramids):
            block_config = level_config[i]  # LIST LIKE [8,8,8,...,8]
            pyramid_residual = OrderedDict()

            # AT THE END OF EACH DCU, WE INCLUDE A CONV(1,1) COMPRESSION LAYER
            # NO NEED FOR THIS AT LEVEL 0
            if i != 0:
                out_planes = num_init_features

                # out_planes = num_init_features = 160
                pyramid_residual['compression_%d' % i] = CompressionBlock(in_planes=num_features, out_planes=out_planes)
                num_features = out_planes

            # serial connect blocks
            # NUM OF ELEMENTS IN block_confis IS THE NUMBER OF DCUs IN THAT PYRAMID
            # CREATING 8 DCUs
            for b, num_layers in enumerate(block_config):
                # FOR EACH DENSELY_CONNECTED_UNIT ***********************************************************
                # num_layers IS ALWAYS 8
                denseblock_params['num_layers'] = num_layers
                denseblock_params['num_input_features'] = num_features  # 160

                # DENSELY_CONNECTED_BLOCK WITH CONV(1,1) INSIDE*********************************************************
                pyramid_residual['residual_denseblock_%d' %(b + 1)] = DenseConnectedBlock(
                    res_factor=res_factor,
                    **denseblock_params)

            # conv before upsampling
            # THIS IS R
            block, num_features = self.create_finalconv(num_features, max_num_feature)

            # CREATING PYRAMID
            pyramid_residual['final_conv'] = block
            self.add_module('pyramid_residual_%d' % (i + 1),
                            nn.Sequential(pyramid_residual))

            # upsample the residual by 2 before reconstruction and next level
            self.add_module(
                'pyramid_residual_%d_residual_upsampler' % (i + 1),
                self.Upsampler(2, num_features))

            # reconstruction convolutions
            reconst_branch = OrderedDict()
            out_channels = num_features
            reconst_branch['final_conv'] = Conv2d(out_channels, 3, 3)
            self.add_module('reconst_%d' % (i + 1),
                            nn.Sequential(reconst_branch))

        init_weights(self)

    # GET V BASED ON PYRAMID-ID
    def get_init_conv(self, idx):
        """choose which init_conv based on curr_scale_idx (1-based)"""
        return getattr(self, 'init_conv_%d' % idx)

    def create_finalconv(self, in_channels, max_channels=None):
        block = OrderedDict()
        if in_channels > max_channels:
            # OUR PATH
            block['final_comp'] = CompressionBlock(in_channels, max_channels)
            block['final_conv'] = Conv2d(max_channels, max_channels, (3, 3))
            out_channels = max_channels
        else:
            block['final_conv'] = Conv2d(in_channels, in_channels, (3, 3))
            out_channels = in_channels
        return nn.Sequential(block), out_channels

    # UPSCALE_FACTOR IS THE CURRENT MODEL SCALE...FROM THE DATASET
    def forward(self, x, upscale_factor=None, blend=1.0):
        if upscale_factor is None:
            upscale_factor = self.max_scale
        else:
            valid_upscale_factors = [
                2 ** (i + 1) for i in range(self.num_pyramids)
            ]
            if upscale_factor not in valid_upscale_factors:
                print("Invalid upscaling factor {}: choose one of: {}".format(upscale_factor, valid_upscale_factors))
                raise SystemExit(1)

        #print("RIKI: gen_dis UPSCALE FACTOR: " + str(upscale_factor))
        # GET THE V FOR THIS UPSCALE FACTOR
        # V- COMPUTATION **********************
        feats = self.get_init_conv(log2(upscale_factor))(x)

        # THIS ENSURES WE ONLY GO DOWN THE RELEVANT PART OF THE PYRAMID
        #print(">>>>>UPSCALE:" +str(upscale_factor))
        for s in range(1, int(log2(upscale_factor)) + 1):

            # PYRAMID- COMPUTATION **********************
            feats = getattr(self, 'pyramid_residual_%d' % s)(feats) + feats

            # UPSAMPLING- COMPUTATION **********************
            feats = getattr(self, 'pyramid_residual_%d_residual_upsampler' % s)(feats)

            # RECONSTRUCTION **********************
            # reconst residual image if reached desired scale /
            # use intermediate as base_img / use blend and s is one step lower than desired scale
            if 2 ** s == upscale_factor or (blend != 1.0 and 2 ** (s + 1) == upscale_factor):
                tmp = getattr(self, 'reconst_%d' % s)(feats)
                # if using blend, upsample the second last feature via bilinear upsampling
                if (blend != 1.0 and s == self.current_scale_idx):
                    #print("HERE: SCALEID:"+str(self.current_scale_idx))
                    base_img = nn.functional.upsample(
                        tmp,
                        scale_factor=2,
                        mode='bilinear',
                        align_corners=True)
                if 2 ** s == upscale_factor:
                    if (blend != 1.0) and s == self.current_scale_idx + 1:
                        tmp = tmp * blend + (1 - blend) * base_img
                    output = tmp

        return output

<file_sep>import gdal
from affine import Affine
import geohash2
import numpy as np
from osgeo import ogr
from osgeo import osr


def retrieve_pixel_value(x,y,data_source):
    #print(data_source.GetProjection())
    #print(data_source.GetGeoTransform())
    forward_transform = Affine.from_gdal(*data_source.GetGeoTransform())
    reverse_transform = ~forward_transform
    px, py = reverse_transform * (x, y)
    px, py = int(px + 0.5), int(py + 0.5)
    pixel_coord = px, py

    return pixel_coord

#RETURNS LAT-LON BOUNDS OF A GEOHASH
def decode_geohash(geocode):
    lat, lon, l1, l2 = geohash2.decode_exactly(geocode)
    return lat-l1, lat+l1, lon-l2, lon+l2

def get_gdal_obj(filename):
    return gdal.Open(filename)

def get_pixel_from_lat_lon(latlons, dataset):
    #converting coordinate systems
    # Setup the source projection
    source = osr.SpatialReference()
    source.ImportFromWkt(dataset.GetProjection())
    #print(source)
    # The target projection
    target = osr.SpatialReference()
    target.ImportFromEPSG(4326)
    # Create the transform - this can be used repeatedly
    transform = osr.CoordinateTransformation(target, source)

    pixels=[]
    for lat,lon in latlons:
        x, y, z = transform.TransformPoint(lon, lat)
        #print("TRANSFORMED:", x, y, z)
        pixels.append(retrieve_pixel_value(x, y, dataset))
    return pixels

# CROP OUT A SMALLER GEOHASH FROM A LARGER GEOHASH TILE
def crop_geohash(geocode, datafile):
    lat1, lat2, lon1, lon2 = decode_geohash(geocode)
    return crop_section(lat1, lat2, lon1, lon2, datafile)

# CROPPING A RECTANGLE OUT OF AN IMAGE
def crop_section(lat1, lat2, lon1, lon2, datafile):
    latlons = []
    latlons.append((lat1, lon1))
    latlons.append((lat2, lon1))
    latlons.append((lat2, lon2))
    latlons.append((lat1, lon2))
    return get_pixel_from_lat_lon(latlons, datafile)


# ************************METHOD CALLS
geocode='9xjm5k'
print("DECODED:",geocode,decode_geohash(geocode))

filename = "/s/chopin/e/proj/sustain/sapmitra/super_resolution/SRImages/train/9xjm5_20190407.tif" #path to raster
lon,lat = -104.9635,40.2769
latlons = []
latlons.append((lat,lon))
print("PIXEL:",get_pixel_from_lat_lon(latlons,get_gdal_obj(filename)))

print("PIXEL:",crop_geohash(geocode,get_gdal_obj(filename)))
<file_sep>import os
import os.path as osp
import glob
import torch.nn as nn
from math import floor
from numpy import random
from PIL import Image
import numpy as np
import torch
from torchvision.utils import save_image, make_grid
from skimage import img_as_float
from skimage.color import rgb2ycbcr
from skimage.measure import compare_psnr, compare_ssim


def get_filenames(source, ext):
    img_paths = []
    for roots, dir, files in os.walk(source):
        for file in files:
            if ext in file:
                file_abs_path = os.path.join(roots, file)
                tokens = file_abs_path.split('/');
                ln = len(tokens)
                geohash = tokens[ln - 2]

                f_tokens = file.split('$')
                date = f_tokens[0]

                img_paths.append(file_abs_path)
    return img_paths


# RETURNS ALL FILENAMES IN THE GIVEN DIRECTORY IN AN ARRAY
def get_filenames1(source, image_format):
    # If image_format is a list
    if source is None:
        return []
    # Seamlessy load single file, list of files and files from directories.
    source_fns = []
    if isinstance(source, str):
        if os.path.isdir(source) or source[-1] == '*':
            if isinstance(image_format, list):
                for fmt in image_format:
                    source_fns += get_filenames(source, fmt)
            else:
                source_fns = sorted(
                    glob.glob("{}/*.{}".format(source, image_format)))
        elif os.path.isfile(source):
            source_fns = [source]
        assert (all([is_image_file(f) for f in source_fns
                     ])), "Given files contain files with unsupported format"
    elif len(source) and isinstance(source[0], str):
        for s in source:
            source_fns.extend(get_filenames(s, image_format=image_format))
    return source_fns


def load_dataset(args, img_format):
    files = {'train':{},'test':{}}

    for phase in ['train','test']:
        for ft in ['source','target']:
            if args[phase].dataset.path[ft]:
                files[phase][ft] = get_filenames(args[phase].dataset.path[ft], image_format=img_format)
            else:
                files[phase][ft] = []

    return files['train'],files['test']

def is_image_file(filename, ext):
    return any(filename.lower().endswith(ext))



def image_loader(path, mode='RGB'):
    # open path as file to avoid ResourceWarning
    # (https://github.com/python-pillow/Pillow/issues/835)
    with open(path, 'rb') as f:
        with Image.open(f) as img:
            return img.convert(mode)


def imageResize(image, pixel_res):
    new_image = image.resize((pixel_res, pixel_res))
    return new_image

# IMAGE DOWNGRADING
    # img: PIL image
    # DECREASE RESOLUTION OF IMAGE BY FACTOR OF 2^pow
def downscaleImage(img, scale, method=Image.BICUBIC):
    ow, oh = img.size

    h = int(round(oh / scale))
    w = int(round(ow / scale))
    return img.resize((w, h), method)


def random_rot90(img, r=None):
    if r is None:
        r = random.random() * 4  # TODO Check and rewrite func
    if r < 1:
        return img.transpose(Image.ROTATE_90)
    elif r < 2:
        return img.transpose(Image.ROTATE_270)
    elif r < 3:
        return img.transpose(Image.ROTATE_180)
    else:
        return img


# RANDOM FLIPPING AND ROTATION OF THE INPUT AND TARGET IMAGES
def augment_pairs(img1, img2):
    vflip = random.random() > 0.5
    hflip = random.random() > 0.5
    rot = random.random() * 4
    if hflip:
        img1 = img1.transpose(Image.FLIP_LEFT_RIGHT)
        img2 = img2.transpose(Image.FLIP_LEFT_RIGHT)
    if vflip:
        img1 = img1.transpose(Image.FLIP_TOP_BOTTOM)
        img2 = img2.transpose(Image.FLIP_TOP_BOTTOM)

    img1 = random_rot90(img1, rot)
    img2 = random_rot90(img2, rot)
    return img1, img2

# COMBINE 3 IMAGES INTO A SINGLE IMAGE-GRID
'''
def combine_images1(lr, hr, gen, scale):
    img_hr = make_grid(hr, nrow=1, normalize=False, padding=10, pad_value=0)

    img_lr = nn.functional.interpolate(lr, scale_factor=scale)

    img_lr = make_grid(img_lr, nrow=1, normalize=False, padding=10, pad_value=0)
    gen_img = make_grid(gen, nrow=1, normalize=False, padding=10, pad_value=0)

    img_grid = torch.cat((img_lr.cpu(), img_hr.cpu(), gen_img.cpu(),), -1)
    return img_grid
'''

# COMBINE 3 IMAGES INTO A SINGLE IMAGE-GRID
def combine_images(lr, hr, gen, scale):

    img_lr = nn.functional.interpolate(lr, scale_factor=scale)
    img_grid = torch.cat((img_lr.detach().cpu(), hr.detach().cpu(), gen.detach().cpu()), 0)
    # nrow= NUMBER OF IMAGES IN A SINGLE ROW
    img = make_grid(img_grid, nrow=3, normalize=False, padding=3, pad_value=255)

    return img

# COMBINE 4 IMAGES INTO A SINGLE IMAGE-GRID...DURING TESTING
def combine_images_test(lr, hr, bic, gen, scale):

    img_lr = nn.functional.interpolate(lr, scale_factor=scale)
    img_grid = torch.cat((img_lr.detach().cpu(), hr.detach().cpu(), bic.detach().cpu(), gen.detach().cpu()), 0)
    # nrow= NUMBER OF IMAGES IN A SINGLE ROW
    img = make_grid(img_grid, nrow=4, normalize=False, padding=3, pad_value=255)

    return img

# Converts a Tensor into a Numpy array
def tensor2im(image_tensor, mean=(0.5, 0.5, 0.5), stddev=2.):
    image_numpy = image_tensor[0].cpu().float().numpy()
    image_numpy = (np.transpose(image_numpy,
                                (1, 2, 0)) * stddev + np.array(mean)) * 255.0
    image_numpy = image_numpy.clip(0, 255)
    return np.around(image_numpy).astype(np.uint8)



def crop_boundaries(im, cs):
    if cs > 1:
        return im[cs:-cs, cs:-cs, ...]
    else:
        return im

def mod_crop(im, scale):
    h, w = im.shape[:2]
    # return im[(h % scale):, (w % scale):, ...]
    return im[:h - (h % scale), :w - (w % scale), ...]


def eval_psnr_and_ssim(im1, im2, scale):
    im1_t = np.atleast_3d(img_as_float(im1))
    im2_t = np.atleast_3d(img_as_float(im2))

    if im1_t.shape[2] == 1 or im2_t.shape[2] == 1:
        im1_t = im1_t[..., 0]
        im2_t = im2_t[..., 0]

    else:
        im1_t = rgb2ycbcr(im1_t)[:, :, 0:1] / 255.0
        im2_t = rgb2ycbcr(im2_t)[:, :, 0:1] / 255.0

    if scale > 1:
        im1_t = mod_crop(im1_t, scale)
        im2_t = mod_crop(im2_t, scale)

        # NOTE conventionally, crop scale+6 pixels (EDSR, VDSR etc)
        im1_t = crop_boundaries(im1_t, int(scale) + 6)
        im2_t = crop_boundaries(im2_t, int(scale) + 6)

    psnr_val = compare_psnr(im1_t, im2_t)
    ssim_val = compare_ssim(
        im1_t,
        im2_t,
        win_size=11,
        gaussian_weights=True,
        multichannel=True,
        data_range=1.0,
        K1=0.01,
        K2=0.03,
        sigma=1.5)

    return psnr_val, ssim_val

<file_sep>from collections import  defaultdict

errors_accum = defaultdict(list)

errors_accum['l1x1'].append(11)
errors_accum['l1x1'].append(13)

print(errors_accum)
