
import os
import shutil
import pandas as pd

LR_source = 'raw_data\Dataset'
HR_source = 'raw_data\Dataset'

LR_dis = 'data\SR_data'
HR_dis = 'data\SR_data'

df = pd.read_csv('data\meta_data.csv')

"""
the flag being true means that the image is not in right shape or most of the pixels is just zeros
the csv file is made in a separated notebook

"""

for index, row in df.iterrows():
    if not row['flag']:
        lr_path_source = os.path.join(LR_source, row['LR_path'])
        hr_path_source = os.path.join(HR_source, row['HR_path'])
        lr_path_dest = os.path.join(LR_dis, row['LR_path'])
        hr_path_dest = os.path.join(HR_dis, row['HR_path'])
        
        shutil.move(lr_path_source, lr_path_dest)
        
        shutil.move(hr_path_source, hr_path_dest)


raw_data/preprocessing.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
import setuptools

with open("README.md", "r", encoding="utf-8") as f:
    long_description = f.read()


__version__ = "0.0.0"

REPO_NAME = "satellite-images-super-resolution-using-EEGAN-in-pytorch-"
AUTHOR_USER_NAME = "zayn309"
SRC_REPO = "SR"
AUTHOR_EMAIL = "joozain309@gmail.com"



setuptools.setup(
    name=SRC_REPO,
    version=__version__,
    author=AUTHOR_USER_NAME,
    author_email=AUTHOR_EMAIL,
    description="A small python package for image super resolution",
    long_description=long_description,
    long_description_content="text/markdown",
    url=f"https://github.com/{AUTHOR_USER_NAME}/{REPO_NAME}",
    project_urls={
        "Bug Tracker": f"https://github.com/{AUTHOR_USER_NAME}/{REPO_NAME}/issues",
    },
    package_dir={"": "src"},
    packages=setuptools.find_packages(where="src")
)
setup.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
import torch.nn as nn
import numpy as np
from abc import abstractmethod


class BaseModel(nn.Module):
    """
    Base class for all models
    """
    @abstractmethod
    def forward(self, *inputs):
        """
        Forward pass logic

        :return: Model output
        """
        raise NotImplementedError

    def __str__(self):
        """
        Model prints with number of trainable parameters
        # """
        # name = None
        model_parameters = filter(lambda p: p.requires_grad, self.parameters())
        params = sum([np.prod(p.size()) for p in model_parameters])
        return str(self.__class__.__name__) + ' that has {} Trainable parameters'.format(params)
src/sr/base/base_model.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
import torch
from abc import abstractmethod
from pathlib import Path
from sr.utils.utils import dict2str

class BaseTrainer:
    """
    Base class for all trainers
    """
    def __init__(self, model_G,model_D , opt_G , opt_D , config , logger , data_loader  , metrics ,  data_loader_valid = None, lr_scheduler=None ):
        
        self.config = config
        self.metrics = metrics
        # setup GPU device if available, move model into configured device
        self.logger = logger
        self.device, device_ids = self._prepare_device(config.n_gpu)
        self.model_G = model_G.to(self.device)
        self.model_D = model_D.to(self.device)
        
        if len(device_ids) > 1:
            self.model_G = torch.nn.DataParallel(model_G, device_ids=device_ids)
            self.model_D = torch.nn.DataParallel(model_D, device_ids=device_ids)
            

        self.opt_G = opt_G
        self.opt_D = opt_D
        

        self.data_loader = data_loader
        self.data_loader_valid = data_loader_valid
        
        self.epochs = self.config.train.niter
        self.current_epoch = 1
        self.save_period = self.config.logger.save_checkpoint_freq
        
        self.checkpoint_dir = Path(config.train.save_dir)
        
        self.lr_scheduler = lr_scheduler

        if config.train.resume is not None:
            self._resume_checkpoint(config.train.resume)

    @abstractmethod
    def _train_epoch(self, epoch):
        """
        Training logic for an epoch

        :param epoch: Current epoch number
        """
        raise NotImplementedError

    def train(self):
        """
        Full training logic
        """
        self.logger.info(f'\n ===> started training <===\n' +
                 f'generator: \n {str(self.model_G)}\n' +
                 f'Discrimininator: \n {str(self.model_D)}\n' +
                 f'optimizers: {self.config.optimizer.type}\n' +
                 f'the training will last for {self.epochs} starting from epoch {self.current_epoch}\n'
                 f'the training will be done on {self.device} device')

        for epoch in range(self.current_epoch, self.epochs + 1):
            print(f'started epoch numner {epoch}')
            result = self._train_epoch(epoch)
            #result = {'psnr': 30.55, 'ssim' : 0.6565 , 'mse': 2.46452, 'vgg Loss': 502.012}
            # save logged informations into log dict
            
            log = {'results of epoch ': epoch}
            log.update(result)
            log = dict2str(log)
            # print logged informations to the screen
            
            self.logger.info(log)

            if epoch % self.save_period == 0:
                self._save_checkpoint(epoch)

    def _prepare_device(self, n_gpu_use): # working 
        """
        setup GPU device if available, move model into configured device
        """
        n_gpu = torch.cuda.device_count()
        if n_gpu_use > 0 and n_gpu == 0:
            self.logger.warning("Warning: There\'s no GPU available on this machine,"
                                "training will be performed on CPU.")
            n_gpu_use = 0
        if n_gpu_use > n_gpu:
            self.logger.warning("Warning: The number of GPU\'s configured to use is {}, but only {} are available "
                                "on this machine.".format(n_gpu_use, n_gpu))
            n_gpu_use = n_gpu
        device = torch.device('cuda:0' if n_gpu_use > 0 else 'cpu')
        list_ids = list(range(n_gpu_use))
        return device, list_ids

    def _save_checkpoint(self, epoch): # working 
        """
        Saving checkpoints

        :param epoch: current epoch number
        :param log: logging information of the epoch
        :param save_best: if True, rename the saved checkpoint to 'model_best.pth'
        """
        arch = type(self.model_G).__name__
        state = {
            'arch': arch,
            'epoch': epoch,
            'state_dict_G': self.model_G.state_dict(),
            'state_dict_D': self.model_D.state_dict(),
            'opt_G': self.opt_G.state_dict(),
            'opt_D': self.opt_D.state_dict(),
            'config': self.config
        }
        filename = str(self.checkpoint_dir / 'checkpoint-epoch{}.pth'.format(epoch))
        torch.save(state, filename)
        self.logger.info("Saving checkpoint: {} ...".format(filename))

    def _resume_checkpoint(self, resume_path): # working
        """
        Resume from saved checkpoints

        :param resume_path: Checkpoint path to be resumed
        """
        resume_path = str(resume_path)
        self.logger.info("Loading checkpoint: {} ...".format(resume_path))
        checkpoint = torch.load(resume_path)
        self.current_epoch = checkpoint['epoch'] + 1

        # load architecture params from checkpoint.
        if checkpoint['config'].arch != self.config.arch:
            self.logger.warning("Warning: Architecture configuration given in config file is different from that of "
                                "checkpoint. This may yield an exception while state_dict is being loaded.")
        self.model_G.load_state_dict(checkpoint['state_dict_G'])
        self.model_D.load_state_dict(checkpoint['state_dict_D'])

        # load optimizer state from checkpoint only when optimizer type is not changed.
        if checkpoint['config'].optimizer.type != self.config.optimizer.type:
            self.logger.warning("Warning: Optimizer type given in config file is different from that of checkpoint. "
                                "Optimizer parameters not being resumed.")
        else:
            self.opt_G.load_state_dict(checkpoint['opt_G'])
            self.opt_D.load_state_dict(checkpoint['opt_D'])

        self.logger.info("Checkpoint loaded. Resume training from epoch {}".format(self.current_epoch))

src/sr/base/base_trainer.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
from torchvision import datasets, utils
from sr.base.base_data_loader import BaseDataLoader
from sr.datsets.sr_dataset import SrDataset
import albumentations as A
from albumentations.pytorch import ToTensorV2
from albumentations import (
    HorizontalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90,
    Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,
    IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, IAAPiecewiseAffine,
    IAASharpen, IAAEmboss, RandomBrightnessContrast, Flip, OneOf, Compose,
    BboxParams, RandomCrop, Normalize, Resize, VerticalFlip
)

from sr.utils.utils import get_config

class MnistDataLoader(BaseDataLoader):
    """
    MNIST data loading demo using BaseDataLoader
    """
    def __init__(self, data_dir, batch_size, shuffle=True, validation_split=0.0, num_workers=0, training=True):
        trsfm = A.Compose([
            A.Normalize((0.1307,), (0.3081,)),
            A.ToTensor(),
        ])
        self.data_dir = data_dir
        self.dataset = datasets.MNIST(self.data_dir, train=training, download=True, transform=trsfm)
        super().__init__(self.dataset, batch_size, shuffle, validation_split, num_workers)

class SR_dataLoader(BaseDataLoader):
    
    def __init__(self,config):
        
        self.config = config
        
        self.batch_size = self.config.data_loader.args.batch_size
        self.data_dir_HR = self.config.data_loader.args.data_dir_HR
        self.data_dir_LR = self.config.data_loader.args.data_dir_LR
        self.shuffle = self.config.data_loader.args.shuffle
        self.validation_split = self.config.data_loader.args.validation_split
        self.num_workers = self.config.data_loader.args.num_workers
        
        lr_data_transforms = A.Compose([
            A.Resize(64,64),
            ToTensorV2(),
        ])
        hr_data_transforms = A.Compose([
            A.Resize(256,256),
            ToTensorV2(),
        ])
        both_transorms = A.Compose([
            A.HorizontalFlip(p=0.5),  # Apply horizontal flip with 50% probability
            A.RandomRotate90(p=0.5),
        ],additional_targets={'target':"image"})
        self.dataset = SrDataset(LR_root=self.data_dir_LR,HR_root=self.data_dir_HR,
                                 lr_transform = lr_data_transforms, hr_transform = hr_data_transforms, both_transorms = both_transorms )
        super().__init__(self.dataset, self.batch_size, self.shuffle, self.validation_split, self.num_workers)
        
def test():
    sr_loader = SR_dataLoader(config= get_config('config.json'))
    data_loader_val = sr_loader.split_validation()
    print(f'the number of training batches: {len(sr_loader)}')
    print(f'the number of validation batches: {len(data_loader_val)}')
        
src/sr/data_loader/data_loaders.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
from torch.utils.data import Dataset
import glob
import torch
import tifffile as tiff
import numpy as np

class SrDataset(Dataset):
    def __init__(self,LR_root , HR_root,lr_transform = None,hr_transform = None ,both_transorms = None):
        
        self.LR_root = LR_root
        self.HR_root = HR_root
        self.lr_transform = lr_transform
        self.hr_transform = hr_transform
        self.both_transorms = both_transorms
        self.LR_images = list(sorted(set(glob.glob(self.LR_root+"*.tif"))))
        self.HR_images = list(sorted(set(glob.glob(self.HR_root+"*.tif"))))
        
        
    def __len__(self):
        return len(self.HR_images)
    
    def _scaleCCC(self,x):
        x = np.where(np.isnan(x), 0, x)
        return((x - np.nanpercentile(x, 2))/(np.nanpercentile(x, 98) - np.nanpercentile(x,2)))
    
    def __getitem__(self, index):
        lr_image_path = self.LR_images[index]
        hr_image_path = self.HR_images[index]

        lr_image = self._scaleCCC(np.array(tiff.imread(lr_image_path), dtype='float32'))
        hr_image = self._scaleCCC(np.array(tiff.imread(hr_image_path), dtype='float32'))
        if self.both_transorms:
            transformed = self.both_transorms(image = lr_image,target = hr_image)
            lr_image = transformed["image"]
            hr_image = transformed["target"]
        if self.lr_transform:
            lr_image = self.lr_transform(image =lr_image)['image']
        if self.hr_transform:
            hr_image = self.hr_transform(image = hr_image)['image']
            
        return lr_image, hr_image

src/sr/datsets/sr_dataset.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
import os 
import sys 
import logging

logging_string = '''[Time: %(asctime)s]
[logging level: %(levelname)s]
[message: %(message)s]
-----------------------------------'''
log_dir = 'logs'
log_filepath = os.path.join(log_dir,"running_log.log")
os.makedirs(log_dir,exist_ok=True)


logging.basicConfig(
    level=logging.INFO,
    format=logging_string,
    
    handlers= [
        logging.FileHandler(log_filepath),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger("SR")

src/sr/logger/__init__.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
import torch
from skimage.metrics import structural_similarity as ssim
import numpy as np
from sr.utils.utils import get_config
import warnings
warnings.filterwarnings("ignore")
class Metrics():
    def __init__(self, config, log = None):
        metrics = config.metrics
        self.log = log
        implemented = ['psnr','mse', 'ssim']
        for value in metrics:
            assert value in implemented , f"sorry the {value} metrics is not imlemented"
        
    def _psnr(self, original, distorted):
        mse = torch.mean((original - distorted) ** 2, dim=(1,2,3))
        max_pixel = torch.max(original)
        return 10 * torch.log10((max_pixel ** 2) / mse)
    
    def _ssim(self, original, distorted):
        original = original.permute(0, 2, 3, 1).cpu().detach().numpy()  # Convert to numpy
        distorted = distorted.permute(0, 2, 3, 1).cpu().detach().numpy()  # Convert to numpy
        ssim_values = []
        for i in range(original.shape[0]):
            data_range = max(np.max(original[i]) , np.max(distorted[i])) - min(np.min(original[i]) , np.min(distorted[i]))
            ssim_values.append(ssim(original[i], distorted[i], multichannel=True, channel_axis = 2, data_range = data_range ))
        return torch.tensor(ssim_values)
    
    def _mse(self, original, distorted):
        mse = torch.mean((original - distorted) ** 2)
        return mse
    
    def result(self, original, distorted):
        psnr_value = self._psnr(original, distorted).mean()
        ssim_value = self._ssim(original, distorted).mean()
        mse_value = self._mse(original, distorted)
        
        return {
            'psnr': psnr_value.item(),
            'ssim': ssim_value.item(),
            'mse': mse_value.item()
        }

# def test():
#     config = get_config('config.json')
#     met = Metrics(config)
#     input_tensor = torch.randn(4,4,256,56)
#     distorted_tensor = torch.randn(4,4,256,56)
#     result = met.result(input_tensor, distorted_tensor)
    
#     print(result)
    
# test()
src/sr/metrics/metrics.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
import torch
import torch.nn as nn
from sr.base.base_model import BaseModel
class Discriminator_VGG_128(BaseModel):
    def __init__(self, in_nc, nf):
        super(Discriminator_VGG_128, self).__init__()
        # [64, 128, 128]
        self.conv0_0 = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
        self.conv0_1 = nn.Conv2d(nf, nf, 4, 2, 1, bias=False)
        self.bn0_1 = nn.BatchNorm2d(nf, affine=True)
        # [64, 64, 64]
        self.conv1_0 = nn.Conv2d(nf, nf * 2, 3, 1, 1, bias=False)
        self.bn1_0 = nn.BatchNorm2d(nf * 2, affine=True)
        self.conv1_1 = nn.Conv2d(nf * 2, nf * 2, 4, 2, 1, bias=False)
        self.bn1_1 = nn.BatchNorm2d(nf * 2, affine=True)
        # [128, 32, 32]
        self.conv2_0 = nn.Conv2d(nf * 2, nf * 4, 3, 1, 1, bias=False)
        self.bn2_0 = nn.BatchNorm2d(nf * 4, affine=True)
        self.conv2_1 = nn.Conv2d(nf * 4, nf * 4, 4, 2, 1, bias=False)
        self.bn2_1 = nn.BatchNorm2d(nf * 4, affine=True)
        # [256, 16, 16]
        self.conv3_0 = nn.Conv2d(nf * 4, nf * 8, 3, 1, 1, bias=False)
        self.bn3_0 = nn.BatchNorm2d(nf * 8, affine=True)
        self.conv3_1 = nn.Conv2d(nf * 8, nf * 8, 4, 2, 1, bias=False)
        self.bn3_1 = nn.BatchNorm2d(nf * 8, affine=True)
        # [512, 8, 8]
        self.conv4_0 = nn.Conv2d(nf * 8, nf * 8, 3, 1, 1, bias=False)
        self.bn4_0 = nn.BatchNorm2d(nf * 8, affine=True)
        self.conv4_1 = nn.Conv2d(nf * 8, nf * 8, 4, 2, 1, bias=False)
        self.bn4_1 = nn.BatchNorm2d(nf * 8, affine=True)

        self.linear1 = nn.Linear(2048 * 4 * 4, 100)
        self.linear2 = nn.Linear(100, 1)

        # activation function
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)

    def forward(self, x):
        fea = self.lrelu(self.conv0_0(x))
        fea = self.lrelu(self.bn0_1(self.conv0_1(fea)))

        fea = self.lrelu(self.bn1_0(self.conv1_0(fea)))
        fea = self.lrelu(self.bn1_1(self.conv1_1(fea)))

        fea = self.lrelu(self.bn2_0(self.conv2_0(fea)))
        fea = self.lrelu(self.bn2_1(self.conv2_1(fea)))

        fea = self.lrelu(self.bn3_0(self.conv3_0(fea)))
        fea = self.lrelu(self.bn3_1(self.conv3_1(fea)))

        fea = self.lrelu(self.bn4_0(self.conv4_0(fea)))
        fea = self.lrelu(self.bn4_1(self.conv4_1(fea)))

        fea = fea.view(fea.size(0), -1)
        fea = self.lrelu(self.linear1(fea))
        out = self.linear2(fea)
        return out
    
    def __str__(self):
        return f'{super().__str__()} '

def test_Discriminator():
  batch_size = 8
  in_channels = 4
  height, width = 256,256
  features = 64
  input_tensor = torch.randn(batch_size, in_channels, height, width)
  model = Discriminator_VGG_128(in_channels,features)
  output = model(input_tensor)
  print(output.shape)
  
src/sr/models/EEGAN/Discriminator.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
from sr.base.base_model import BaseModel
import torch
import torch.nn as nn
import kornia
from sr.models.EEGAN.Gen_components import  BeginEdgeConv,EESN_dense,MaskConv,FinalConv

class EESN(BaseModel):
    def __init__(self, scale_factor):
        super(EESN, self).__init__()
        self.beginEdgeConv = BeginEdgeConv() #  Output 64*64*64 input 3*64*64
        self.denseNet = EESN_dense() 
        self.maskConv = MaskConv() # Output 256*64*64
        self.finalConv = FinalConv(256,64,scale_factor) # Output 4*256*256

    def forward(self, x):
        x_lap = kornia.filters.laplacian(x, 3) # see kornia laplacian kernel
        x1 = self.beginEdgeConv(x_lap)
        x2 = self.denseNet(x1)
        x3 = self.maskConv(x1)
        x4 = x3*x2 + x2
        x_learned_lap = self.finalConv(x4)

        return x_learned_lap ,x_lap

src/sr/models/EEGAN/EESN.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
from sr.base.base_model import BaseModel
import torch
import torch.nn as nn
from functools import partial
import functools
import torch.nn.functional as F
from kornia.filters import laplacian
from sr.utils.utils import (initialize_weights, make_layer)


class UpsamplingBlock(nn.Module):
    def __init__(self, in_channels, out_channels, scale_factor):
        super(UpsamplingBlock, self).__init__()
        self.conv_prelu = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.PReLU(out_channels)
        )
        self.sub_pixel_conv = nn.PixelShuffle(scale_factor)

    def forward(self, x):
        x = self.conv_prelu(x)
        x = self.sub_pixel_conv(x)
        return x


class ResidualDenseBlock_5C(nn.Module):
    def __init__(self, nf=64, gc=32, bias=True):
        super(ResidualDenseBlock_5C, self).__init__()
        # gc: growth channel, i.e. intermediate channels
        self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1, bias=bias)
        self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1, bias=bias)
        self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1, bias=bias)
        self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1, bias=bias)
        self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1, bias=bias)
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)

        # initialization
        initialize_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)

    def forward(self, x):
        x1 = self.lrelu(self.conv1(x))
        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
        return x5 * 0.2 + x
    
class RRDB(nn.Module):
    '''Residual in Residual Dense Block'''

    def __init__(self, nf, gc=32):
        super(RRDB, self).__init__()
        self.RDB1 = ResidualDenseBlock_5C(nf, gc)
        self.RDB2 = ResidualDenseBlock_5C(nf, gc)
        self.RDB3 = ResidualDenseBlock_5C(nf, gc)

    def forward(self, x):
        out = self.RDB1(x)
        out = self.RDB2(out)
        out = self.RDB3(out)
        return out * 0.2 + x
    
class RRDBNet(nn.Module):
    def __init__(self, in_nc, out_nc, nf, nb, gc=32):
        super(RRDBNet, self).__init__()
        RRDB_block_f = functools.partial(RRDB, nf=nf, gc=gc)

        self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
        self.RRDB_trunk = make_layer(RRDB_block_f, nb)
        self.trunk_conv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
        #### upsampling
        self.upconv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
        self.upconv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
        self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
        self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
        # self.up = UpsamplingBlock(nf,nf,4)
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)

    def forward(self, x):
        fea = self.conv_first(x)
        trunk = self.trunk_conv(self.RRDB_trunk(fea))
        fea = fea + trunk

        fea = self.lrelu(self.upconv1(F.interpolate(fea, scale_factor=2, mode='nearest')))
        fea = self.lrelu(self.upconv2(F.interpolate(fea, scale_factor=2, mode='nearest')))
        out = self.conv_last(self.lrelu(self.HRconv(fea)))
        # out = self.lrelu(self.upconv1(fea))
        # out = self.lrelu(self.upconv2(fea))
        # out = self.lrelu(self.HRconv(fea))
        # out = self.up(fea)

        return out

# UDSN components
class ConvLayer(nn.Module):  
    def __init__(self, in_channels,out_channels, kernel_size, strides=[1, 1, 1], padding=1): 
        assert len(strides) == 3
        super(ConvLayer, self).__init__()
        self.conv1 = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=strides[0],
            padding=padding
        )
        self.conv2 = nn.Conv2d(
            in_channels=in_channels,  
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=strides[1],
            padding=padding
        )
        self.conv3 = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=strides[2],
            padding=padding
        )
        
        self.activation = nn.LeakyReLU(0.2)
    
    def forward(self, x, act=True ):
        if isinstance(x, list):
            if act:
                return [
                    self.activation(self.conv1(x[0])),
                    self.activation(self.conv2(x[1])),
                    self.activation(self.conv3(x[2]))
                ]
            else:
                return [
                    self.conv1(x[0]),
                    self.conv2(x[1]),
                    self.conv3(x[2])
                ]
        else:
            if act:
                return [
                    self.activation(self.conv1(x)),
                    self.activation(self.conv2(x)),
                    self.activation(self.conv3(x))
                ]
            else:
                return [
                    self.conv1(x),
                    self.conv2(x),
                    self.conv3(x)
                ]

class Basic_block(nn.Module):
    def __init__(self, in_channels, features):
        super(Basic_block, self).__init__()
        
        self.conv1 = ConvLayer(in_channels ,features ,kernel_size= 3)
        
        self.activation = nn.LeakyReLU(0.2)
        
        self.conv2 = ConvLayer(features*3, in_channels, kernel_size = 1, padding = 0)
        
        
    def forward(self,x):
        
        residual = x
        
        # conv f3
        out = self.conv1.forward(x)
        
        # concatenation
        out = torch.cat(out,dim = 1)
        
        # conv f1
        out = self.conv2.forward(out,act = False)
        
        if (residual is not None):
            if isinstance(residual, list):
                out[0] = self.activation(out[0] + residual[0])
                out[1] = self.activation(out[1] + residual[1])
                out[2] = self.activation(out[2] + residual[2])
            else:
                out[0] = self.activation(out[0] + residual)
                out[1] = self.activation(out[1] + residual)
                out[2] = self.activation(out[2] + residual)
        
        return out
    
class Dense_block(nn.Module):
    def __init__(self,in_channels, features):
        super(Dense_block,self).__init__()
        
        self.block1 = nn.Sequential(
                            nn.Conv2d(
                                in_channels = in_channels,
                                out_channels = features,
                                kernel_size = 3,
                                stride = 1,
                                padding = 1
                            ),
                            nn.LeakyReLU(0.2)
                        )
        self.block2 = Basic_block(features,features)
        self.block3 = Basic_block(features,features)
        self.block4 = Basic_block(features,features)
        self.block5 =  nn.Sequential(
                            nn.Conv2d(
                                in_channels = features*3,
                                out_channels = features,
                                kernel_size = 3,
                                stride = 1,
                                padding = 1
                            ),
                            nn.LeakyReLU(0.2)
                        )
        
    def forward(self,x):
        x = self.block1(x)
        residual = x
        x = self.block2(x)
        x = self.block3(x)
        x = self.block4(x)
        x = torch.cat(x,dim = 1)
        x = self.block5(x)
        return x + residual
  
# EESN components
class BeginEdgeConv(nn.Module):
    def __init__(self):
        super(BeginEdgeConv, self).__init__()

        self.conv_layer1 = nn.Conv2d(4, 64, 3, 1, 1, bias=True)
        self.conv_layer2 = nn.Conv2d(64, 64, 3, 1, 1, bias=True)
        self.conv_layer3 = nn.Conv2d(64, 128, 3, 2, 1, bias=True)
        self.conv_layer4 = nn.Conv2d(128, 128, 3, 1, 1, bias=True)
        self.conv_layer5 = nn.Conv2d(128, 256, 3, 2, 1, bias=True)
        self.conv_layer6 = nn.Conv2d(256, 64, 3, 1, 1, bias=True)
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)

        # initialization
        initialize_weights([self.conv_layer1, self.conv_layer2, self.conv_layer3,
                            self.conv_layer4, self.conv_layer5, self.conv_layer6], 0.1)
        

    def forward(self, x):
      x = self.lrelu(self.conv_layer1(x))
      x = self.lrelu(self.conv_layer2(x))
      x = self.lrelu(self.conv_layer3(x))
      x = self.lrelu(self.conv_layer4(x))
      x = self.lrelu(self.conv_layer5(x))
      x = self.lrelu(self.conv_layer6(x))

      return x

class EESNRRDBNet(nn.Module):
    def __init__(self, in_nc, out_nc, nf, nb, gc=32):
        super(EESNRRDBNet, self).__init__()
        RRDB_block_f = functools.partial(RRDB, nf=nf, gc=gc)

        self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
        self.RRDB_trunk = make_layer(RRDB_block_f, nb)
        self.trunk_conv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
        #### upsampling
        self.upconv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
        self.upconv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
        self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
        self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)

        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)

    def forward(self, x):
        fea = self.conv_first(x)
        trunk = self.trunk_conv(self.RRDB_trunk(fea))
        fea = fea + trunk

        out = self.lrelu(self.conv_last(self.lrelu(self.HRconv(fea))))

        return out

class MaskConv(nn.Module): ## mask branch
    def __init__(self):
        super(MaskConv, self).__init__()

        self.conv_layer1 = nn.Conv2d(64, 64, 3, 1, 1, bias=True)
        self.conv_layer2 = nn.Conv2d(64, 128, 3, 1, 1, bias=True)
        self.conv_layer3 = nn.Conv2d(128, 256, 3, 1, 1, bias=True)
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)

        initialize_weights([self.conv_layer1, self.conv_layer2, self.conv_layer3], 0.1)

    def forward(self, x):
        x = self.lrelu(self.conv_layer1(x))
        x = self.lrelu(self.conv_layer2(x))
        x = self.lrelu(self.conv_layer3(x))
        x = torch.sigmoid(x)

        return x

class FinalConv(nn.Module):
    def __init__(self):
        super(FinalConv, self).__init__()

        self.upconv1 = nn.Conv2d(256, 128, 3, 1, 1, bias=True)
        self.upconv2 = nn.Conv2d(128, 128, 3, 1, 1, bias=True)
        self.HRconv = nn.Conv2d(128, 64, 3, 1, 1, bias=True)
        self.conv_last = nn.Conv2d(64, 4, 3, 1, 1, bias=True)
        # self.up = UpsamplingBlock(64,64,4)
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        

    def forward(self, x):
        x = self.lrelu(self.upconv1(F.interpolate(x, scale_factor=2, mode='nearest')))
        x = self.lrelu(self.upconv2(F.interpolate(x, scale_factor=2, mode='nearest')))
        x = self.conv_last(self.lrelu(self.HRconv(x)))
        # x = self.lrelu(self.upconv1(x))
        # x = self.lrelu(self.upconv2(x))
        # x = self.lrelu(self.HRconv(x))
        # x = self.up(x)
        return x

class EESN(nn.Module):
    def __init__(self):
        super(EESN, self).__init__()
        self.beginEdgeConv = BeginEdgeConv() #  Output 64*64*64 input 3*64*64
        self.denseNet = EESNRRDBNet(64, 256, 64, 5) # RRDB densenet with 64 in kernel, 256 out kernel and 64 intermediate kernel, output: 256*64*64
        self.maskConv = MaskConv() # Output 256*64*64
        self.finalConv = FinalConv() # Output 3*256*256

    def forward(self, x):
        x_lap = laplacian(x, 3) # see kornia laplacian kernel
        x1 = self.beginEdgeConv(x_lap)
        x2 = self.denseNet(x1)
        x3 = self.maskConv(x1)
        x4 = x3*x2 + x2
        x_learned_lap = self.finalConv(x4)

        return x_learned_lap, x_lap
    

class ESRGAN_EESN(BaseModel):
    def __init__(self, in_nc, out_nc, nf, nb):
        super(ESRGAN_EESN, self).__init__()
        self.netRG = RRDBNet(in_nc, out_nc, nf, nb)
        self.netE = EESN()

    def forward(self, x):
        x_base = self.netRG(x) # add bicubic according to the implementation by author but not stated in the paper
        x5, x_lap = self.netE(x_base) # EESN net
        x_sr = x5 + x_base - x_lap

        return x_base, x_sr, x5, x_lap
    def __str__(self):
        return f'{super().__str__()} '
src/sr/models/EEGAN/Gen_components.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
import torch
import torch.nn as nn
from sr.models.EEGAN.Gen_components import Dense_block, UpsamplingBlock


class UDSN(nn.Module):
    def __init__(self,in_channels,features, scale_factor):
        super(UDSN, self).__init__()
        self.dense1 = Dense_block(in_channels, features)
        self.dense2 = Dense_block(features,features)
        self.dense3 = Dense_block(features,features)
        self.dense4 = Dense_block(features,features)
        self.dense5 = Dense_block(features,features)
        self.dense6 = Dense_block(features,features)
        self.up_scale = UpsamplingBlock(features, features ,scale_factor )
        
    def forward(self, x):
        out1 = self.dense1(x)
        out2 = self.dense2(out1) + out1
        out3 = self.dense3(out2) + out1 + out2  # Add out1 and out2
        out4 = self.dense4(out3) + out1 + out2 + out3  # Add out1, out2, and out3
        out5 = self.dense5(out4) + out1 + out2 + out3 + out4  # Add out1, out2, out3, and out4
        out6 = self.dense6(out5) + out1 + out2 + out3 + out4 + out5  # Add out1, out2, out3, out4, and out5
        up_scaled_image = self.up_scale(out6)
        
        return up_scaled_image
src/sr/models/EEGAN/UDSN.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
import torch
import torch.nn as nn 
from sr.models.EEGAN.UDSN import UDSN
from sr.models.EEGAN.EESN import EESN
from sr.base.base_model import BaseModel

class EEGAN_generator(BaseModel):
  def __init__(self, in_nc,features, scale_factor):
    super(EEGAN_generator, self).__init__()
    self.UDSN = UDSN(in_channels=in_nc,features=features,scale_factor=scale_factor)
    self.EESN = EESN(scale_factor=scale_factor)

  def forward(self, x):
    I_base = self.UDSN(x) 
    I_learned_lap ,I_lap = self.EESN(I_base) # EESN net
    I_sr = I_learned_lap + I_base - I_lap

    return I_base, I_sr, I_learned_lap, I_lap
  
def test_generator():
  batch_size = 8
  in_channels = 4
  height, width = 64,64
  features = 64
  input_tensor = torch.randn(batch_size, in_channels, height, width)
  model = EEGAN_generator(in_channels,features)
  A , B, C, D = model(input_tensor)
  print(A.shape ,B.shape,C.shape,D.shape)
  
src/sr/models/EEGAN/generator.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
import torch.nn as nn
import torch
from sr.utils.utils import apply_pca

class CharbonnierLoss(nn.Module):
    """Charbonnier Loss (L1)"""

    def __init__(self, eps=1e-3):
        super(CharbonnierLoss, self).__init__()
        self.eps = eps

    def forward(self, x, y):
        diff = x - y
        loss = torch.mean(torch.sqrt(diff * diff + self.eps ** 2))
        return loss


class ContentLoss(nn.Module):
    def __init__(self,DEVICE, logger):
        super(ContentLoss, self).__init__()
        self.DEVICE = DEVICE
        self.vgg = None
        self.logger = logger
        try:
            self.logger.info("===> loading the vgg19 model <===")
            self.vgg = torch.hub.load('pytorch/vision:v0.9.0', 'vgg19', pretrained=True).features[:36].eval().to(self.DEVICE)
            self.vgg = self.vgg.double()
            self.logger.info("===> the vgg19 loaded successfully <===")
        except :
            self.logger.exception("the vgg model is not loaded properly, cheack you internet connection!")
        
        self.loss = CharbonnierLoss().eval().to(self.DEVICE)
        
        for param in self.vgg.parameters():
            param.requires_grad = False

    def forward(self, IBase, IHR):
        I_Base_compressed = apply_pca(IBase,self.DEVICE)
        IHR_compressed = apply_pca(IHR,self.DEVICE)
        IBase_features = self.vgg(I_Base_compressed)
        IHR_features = self.vgg(IHR_compressed)
        loss = self.loss(IBase_features,IHR_features)
        return loss
    
src/sr/models/EEGAN/loss.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
from torch.optim.lr_scheduler import _LRScheduler

class LR_Scheduler(_LRScheduler):
    def __init__(self, optimizers: list,config):
        self.config = config
        self.initial_lr = self.config.LR_Scheduler.initial_lr
        self.final_lr = self.config.LR_Scheduler.final_lr
        self.decay_factor = self.config.LR_Scheduler.decay_factor
        self.current_lr = self.initial_lr
        self.optimizers = optimizers
        super(LR_Scheduler, self).__init__(optimizers[0])

    def get_lr(self):
        return [self.current_lr] * len(self.optimizers[0].param_groups)

    def step(self, epoch=None):
        if epoch is None:
            epoch = self.last_epoch + 1
        self.last_epoch = epoch

        if self.current_lr > self.final_lr:
            self.current_lr *= self.decay_factor
        else:
            self.current_lr = self.final_lr

        for optimizer in self.optimizers:
            for param_group in optimizer.param_groups:
                param_group['lr'] = self.current_lr

src/sr/models/EEGAN/lr_scheduler.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
import torch
import torch.nn as nn
from tqdm import tqdm 
from sr.base.base_trainer import BaseTrainer
from sr.models.EEGAN.Gen_components import ESRGAN_EESN
from sr.models.EEGAN.Discriminator import Discriminator_VGG_128
from sr.data_loader.data_loaders import SR_dataLoader
from sr.metrics.metrics import Metrics
from sr.models.EEGAN.loss import (CharbonnierLoss, ContentLoss)
from sr.models.EEGAN.lr_scheduler import LR_Scheduler
from sr.utils.utils import plot_examples

class EEGAN_Trainer(BaseTrainer):
    """
    Trainer class
    """
    def __init__(self, config ,logger):


        generator = ESRGAN_EESN(config.network_G.in_nc,config.network_G.out_nc,config.network_G.features,config.network_G.nb)
        discriminator = Discriminator_VGG_128(config.network_D.in_nc,config.network_D.features)
        
        
        opt_G = torch.optim.Adam(generator.parameters(),
                                      config.optimizer.args.lr_G, betas=(config.optimizer.args.beta1_G,config.optimizer.args.beta2_G))
        opt_D = torch.optim.Adam(discriminator.parameters(),
                                      config.optimizer.args.lr_D, betas=(config.optimizer.args.beta1_D,config.optimizer.args.beta2_D))
        
        data_loader = SR_dataLoader(config)
        
        val_data_loader = None
        if config.data_loader.args.validation_split:
            val_data_loader = data_loader.split_validation()
        
        metrics = Metrics(config)
        
        lr_scheduler = LR_Scheduler(optimizers=[opt_G, opt_D],config=config)
        
        super().__init__(generator,discriminator,opt_G,opt_D,config,
                         logger,data_loader,metrics, val_data_loader,lr_scheduler=lr_scheduler)
        
        self.len_epoch = len(self.data_loader)
        
        self.plot_freq = self.config.train.plote_examples_freq
        
        self.G_scaler = torch.cuda.amp.GradScaler()
        self.D_scaler = torch.cuda.amp.GradScaler()
        
        # Losses
        self.loss_terms = self.config.loss.terms
        self.alpha = self.config.loss.alpha
        self.lamda = self.config.loss.lamda
        self.gamma = self.config.loss.gamma
        
        self.consistencyLoss = CharbonnierLoss().to(self.device)
        self.contentLoss = ContentLoss(self.device,logger=self.logger)
        self.BCE = nn.BCEWithLogitsLoss()
        self.l1_loss = nn.L1Loss()
        if self.config.train.use_valid_pretrained_weights:
          self.load_pretrained()
        
    def _train_epoch(self, epoch):
        # L(θG, θD) = Losscont(θG) + αLossadv(θG, θD)+λLosscst(θG)
        self.model_G.train()
        self.model_D.train()
        total_loss = 0
        log = {}
        for batch_idx, (data, target) in enumerate(tqdm(self.data_loader)):
            data, target = data.to(self.device), target.to(self.device)
            
            
            # train the generator with the consistency loss and the content loss
            
            with torch.cuda.amp.autocast():
                I_base , I_sr,_,_ = self.model_G(data)
                disc_fake = self.model_D(I_base)
                # pixel loss
                pixel_loss =  self.l1_loss(I_base,target)
                # vgg loss
                content_loss =  self.contentLoss(I_base,target)
                # adverserial loss
                gen_adv_loss = self.alpha * self.BCE(disc_fake,torch.ones_like(disc_fake))
                # charpo_ loss
                consistency_loss = self.consistencyLoss(I_sr, target)
                loss_gen = content_loss + self.lamda * consistency_loss + self.alpha * gen_adv_loss + self.gamma * pixel_loss
            
            self.opt_G.zero_grad()
            self.G_scaler.scale(loss_gen).backward()
            self.G_scaler.step(self.opt_G)
            self.G_scaler.update()
            
            
            with torch.cuda.amp.autocast():
                
                disc_real = self.model_D(target)
                disc_fake = self.model_D(I_base.detach())
                disc_loss_real = self.BCE(
                    disc_real, torch.ones_like(disc_real) - 0.1 * torch.rand_like(disc_real)
                )
                disc_loss_fake = self.BCE(disc_fake, torch.zeros_like(disc_fake))
                loss_disc = (disc_loss_fake + disc_loss_real) / 2

            self.opt_D.zero_grad()
            self.D_scaler.scale(loss_disc).backward()
            self.D_scaler.step(self.opt_D)
            self.D_scaler.update()
            

            total_loss = loss_disc + loss_gen
            
            log.update({"total_loss" : total_loss.item(),
                        "vgg_loss" : content_loss.item(),
                        "cst_loss" : consistency_loss.item(),
                        "pixel_loss" : pixel_loss.item(),
                        "loss_adv": gen_adv_loss.item(),
                        "loss_disc" : loss_disc.item()
                        })
            
            if batch_idx == self.len_epoch:
                break
        
        
        if epoch % self.config.train.val_freq == 0:
            val_log = self._valid_epoch(epoch,train=False, log=log)
            log.update(val_log)
        try:
          if epoch % self.config.train.train_val_freq == 0:
              val_log = self._valid_epoch(epoch, train = True, log=log)
              log.update(val_log)
          if self.lr_scheduler is not None and epoch >15:
              self.lr_scheduler.step()
          if epoch % self.plot_freq == 0:
              self.logger.info("==> plotting some examples <==")
              self.plot_examples()
        except Exception as e:
          print(e)
          pass
        
        return log

    def _valid_epoch(self, epoch , train = False, log = None):
        """
        Validate after training an epoch

        :param epoch: Integer, current training epoch.
        :return: A log that contains information about validation
        """
        self.model_G.eval()
        
        if log is not None:
          log.update({key: 0 for key in self.config.metrics})
        
        with torch.no_grad():
            for batch_idx, (data, target) in enumerate(self.data_loader_valid) if not train else  enumerate(self.data_loader):
                data, target = data.to(self.device), target.to(self.device)
                _,output,_,_ = self.model_G(data)
                results = self.metrics.result(target,output)
                for key in results:
                    log[key] = log[key] + results[key]
        
        for key in self.config.metrics:
            if train:
                log[key] = log[key] / len(self.data_loader)
            else:
                log[key] = log[key] / len(self.data_loader_valid) 
        if train:
          log.update({"validation on":"training set"})
        else:
          log.update({"validation on":"val set"})
        return log

    def load_pretrained(self):
      path_G = self.config.train.pretrained_G_path
      path_D = self.config.train.pretrained_D_path
      self.logger.info("loading the pretrained weights")
      # load the generator
      weights_G = torch.load(path_G)
      for name, param in self.model_G.named_parameters():
        if name in ['netRG.conv_first.weight', 'netRG.conv_first.bias','netRG.conv_last.weight','netRG.conv_last.bias','netE.beginEdgeConv.conv_layer1.weight','netE.beginEdgeConv.conv_layer1.bias','netE.finalConv.conv_last.weight','netE.finalConv.conv_last.bias']:
          continue
        else:
          try:
            param.data.copy_(weights_G[name])
          except Exception as e:
            print(e)
            print(name)
            break
      # load the disc
      weights_D = torch.load(path_D)
      for name, param in self.model_D.named_parameters():
          if name in ['conv0_0.weight','conv0_0.bias']:
            continue
          else:
            try:
              param.data.copy_(weights_D[name])
            except Exception as e:
              print(e)
              print(name)
              break
      self.logger.info("the pretrained weights was loaded successfully")
    
    def plot_examples(self):
        data, target = next(iter(self.data_loader))
        data = data.to(self.device)
        self.model_G.eval()
        I_base, I_sr, I_learned_lap, I_lap = self.model_G(data)
        plot_examples(I_base,I_lap,I_learned_lap,I_sr,data, target,self.config)
        
src/sr/trainer/EEGAN_trainer.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
import json
from pathlib import Path
from collections import OrderedDict
from datetime import datetime
import os
from box import ConfigBox
import random
import tifffile as tiff
import torch
import numpy as np
from sklearn.decomposition import PCA
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F

def read_json(fname):
    fname = Path(fname)
    with fname.open('rt') as handle:
        return json.load(handle, object_hook=OrderedDict)

def write_json(content, fname):
    fname = Path(fname)
    with fname.open('wt') as handle:
        json.dump(content, handle, indent=4, sort_keys=False)
        
def get_timestamp():
    return datetime.now().strftime('%y%m%d-%H%M%S')


def mkdir(path):
    if not os.path.exists(path):
        os.makedirs(path)


def mkdirs(paths):
    if isinstance(paths, str):
        mkdir(paths)
    else:
        for path in paths:
            mkdir(path)


def get_config(config_path):
    content = read_json(config_path)
    config = ConfigBox(content)
    return config 



def set_random_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

def dict2str(opt, indent_l=1):
    '''dict to string for logger'''
    msg = ''
    for k, v in opt.items():
        if isinstance(v, dict):
            msg += ' ' * (indent_l * 2) + k + ':[\n'
            msg += dict2str(v, indent_l + 1)
            msg += ' ' * (indent_l * 2) + ']\n'
        else:
            msg += ' ' * (indent_l * 2) + k + ': ' + str(v) + '\n'
    return msg


def plot_examples(I_Base, lap, learned_lap, I_sr, IHR, ILR, config):
    def save_images(image_tensor, save_dir, prefix, counter):
        for i, image in enumerate(image_tensor):
            image = image.permute(1, 2, 0).cpu().detach().numpy()
            save_path = str(Path(save_dir) / f"{prefix}_{counter}.tiff")
            tiff.imsave(save_path, image)

    if not isinstance(I_Base, torch.Tensor):
        raise ValueError("Input batch should be a torch.Tensor")

    if len(I_Base.shape) != 4:
        raise ValueError("Input batch should be in the shape of (batch_size, channels, width, height)")

    counter = len(os.listdir(Path(config.examples_dir.base_image))) + 1

    # Save I_Base images
    save_images(I_Base, config.examples_dir.base_image, "I_Base", counter)

    # Save lap images
    save_images(lap, config.examples_dir.edge, "lap", counter)

    # Save learned_lap images
    save_images(learned_lap, config.examples_dir.learned_edge, "learned_lap", counter)

    # Save I_sr images
    save_images(I_sr, config.examples_dir.SR_image, "I_sr", counter)

    # Save IHR images
    save_images(IHR, config.examples_dir.IHR, "IHR", counter)

    # Save ILR images
    save_images(ILR, config.examples_dir.ILR, "ILR", counter)

    # Return the updated counter
    return counter



def apply_pca(images,DEVICE):
    """
    Apply PCA to compress the number of channels from 4 to 3.

    Args:
        images (torch.Tensor): Tensor of images with shape (batch_size, channels, width, height).

    Returns:
        torch.Tensor: Compressed images with shape (batch_size, 3, width, height).
    """
    batch_size, channels, width, height = images.shape

    # Convert to numpy array for PCA
    images_np = images.detach().cpu().numpy()

    # Reshape images to (batch_size, channels, -1)
    images_reshaped = np.reshape(images_np, (batch_size, channels, -1))

    # Reshape to (batch_size, -1, channels) for PCA
    images_for_pca = np.transpose(images_reshaped, (0, 2, 1))

    # Reshape back to (batch_size * width * height, channels) for PCA
    images_for_pca = np.reshape(images_for_pca, (-1, channels))

    # Apply PCA
    pca = PCA(n_components=3)
    compressed_images = pca.fit_transform(images_for_pca)

    # Reshape back to (batch_size, width, height, 3)
    compressed_images = np.reshape(compressed_images, (batch_size, width, height, 3))

    # Reshape to (batch_size, 3, width, height)
    compressed_images = np.transpose(compressed_images, (0, 3, 1, 2))

    # Convert back to torch.Tensor
    compressed_images = torch.from_numpy(compressed_images).to(DEVICE).double()

    return compressed_images


def initialize_weights(net_l, scale=1):
    if not isinstance(net_l, list):
        net_l = [net_l]
    for net in net_l:
        for m in net.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, a=0, mode='fan_in')
                m.weight.data *= scale  # for residual block
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                init.kaiming_normal_(m.weight, a=0, mode='fan_in')
                m.weight.data *= scale
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias.data, 0.0)

def make_layer(block, n_layers):
    layers = []
    for _ in range(n_layers):
        layers.append(block())
    return nn.Sequential(*layers)
src/sr/utils/utils.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
import os
from pathlib import Path
import logging

logging.basicConfig(level=logging.INFO, format='[%(asctime)s]: %(message)s:')


project_name = "sr"

list_of_files = [
    ".github/workflows/.gitkeep",
    f"src/__init__.py",
    f"src/{project_name}/base/__init__.py",
    f"src/{project_name}/utils/__init__.py",
    f"src/{project_name}/logger/__init__.py",
    f"src/{project_name}/model/__init__.py",
    f"src/{project_name}/data_loader/__init__.py",
    f"src/{project_name}/test/__init__.py",
    f"src/{project_name}/trainer/__init__.py",
    f"src/{project_name}/datsets/__init__.py",
    "__init__.py",
    "config.json",
    "config_GAN.json",
    "requirements.txt",
    "setup.py",
    "saved/"
    
]


for filepath in list_of_files:
    filepath = Path(filepath)
    filedir, filename = os.path.split(filepath)

    if filedir != "":
        os.makedirs(filedir, exist_ok=True)
        logging.info(f"Creating directory:{filedir} for the file {filename}")

    
    if (not os.path.exists(filepath)) or (os.path.getsize(filepath) == 0):
        with open(filepath,'w') as f:
            pass
            logging.info(f"Creating empty file: {filepath}")


    
    else:
        logging.info(f"{filename} is already exists")

template.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
# from sr.data_loader.data_loaders import SR_dataLoader
# import tifffile as tiff
from sr.utils.utils import (get_config, plot_examples)
# import os
# from pathlib import Path
# import torch
# import numpy as np
# import matplotlib.pyplot as plt 
# from tqdm import tqdm 
# print(mp.cpu_count())
# class ExampleModel(nn.Module):
#     def __init__(self):
#         super(ExampleModel, self).__init__()
#         self.fc = nn.Linear(10, 10)
#     def __str__(self) -> str:
#         return self.__class__.__name__
#     def forward(self, x):
#         return self.fc(x)
    
# model = ExampleModel()
# opt_G = Adam(model.parameters(), lr=0.001)
# opt_D = Adam(model.parameters(), lr=0.001)
# config = get_config('config.json')

# trainer = BaseTrainer(model=ExampleModel, opt_G=opt_G,opt_D=opt_D, config=config,logger=logger)

# # trainer.save_checkpoint(1)
# #trainer.resume_checkpoint("saved\check points\checkpoint-epoch1.pth")


import matplotlib.pyplot as plt 
import tifffile as tiff 
import numpy as np
import os

# hr = np.array(torch.permute(target,dims = (0,2,3,1))[0,:,:,:])

# plt.imshow(hr)
# plt.show()
# import numpy as np
# from sklearn.decomposition import PCA

# def apply_pca(images):
#     """
#     Apply PCA to compress the number of channels from 4 to 3.

#     Args:
#         images (numpy.ndarray): Tensor of images with shape (batch_size, channels, width, height).

#     Returns:
#         numpy.ndarray: Compressed images with shape (batch_size, 3, width, height).
#     """
#     batch_size, channels, width, height = images.shape

#     # Reshape images to (batch_size, channels, -1)
#     images_reshaped = np.reshape(images, (batch_size, channels, -1))

#     # Reshape to (batch_size, -1, channels) for PCA
#     images_for_pca = np.transpose(images_reshaped, (0, 2, 1))

#     # Reshape back to (batch_size * width * height, channels) for PCA
#     images_for_pca = np.reshape(images_for_pca, (-1, channels))

#     # Apply PCA
#     pca = PCA(n_components=3)
#     compressed_images = pca.fit_transform(images_for_pca)

#     # Reshape back to (batch_size, width, height, 3)
#     compressed_images = np.reshape(compressed_images, (batch_size, width, height, 3))

#     # Reshape to (batch_size, 3, width, height)
#     compressed_images = np.transpose(compressed_images, (0, 3, 1, 2))
    
#     compressed_images = torch.from_numpy(compressed_images)

#     return compressed_images


# import torch
# from sr.data_loader.data_loaders import SR_dataLoader
# config = get_config("config.json")

# data_loader = SR_dataLoader(config)
# data,target = next(iter(data_loader))
# compressed = apply_pca(target)
# # print()
# plt.imshow(compressed[0,:,:,:])
# plt.show()
bad_data_lr = ['LR.33.tif',
              'LR.60.tif',
              'LR.72.tif',
              'LR.86.tif',
              'LR.181.tif',
              'LR.268.tif']

bad_data_hr = [ 'HR.33.tif',
                'HR.60.tif',
                'HR.72.tif',
                'HR.86.tif',
                'HR.181.tif',
                'HR.268.tif']

# Define the directory paths
hr_dir = "data/SR_data/HR"
lr_dir = "data/SR_data/LR"

# Function to delete files with specified numbers
def delete_files_with_numbers(directory, bad_data_lr):
    for file_ in bad_data_lr:

        file_path = os.path.join(directory, file_)
        if os.path.exists(file_path):
            os.remove(file_path)
            print(f"Deleted {file_}")

# Delete bad LR files
delete_files_with_numbers(lr_dir, bad_data_lr)
delete_files_with_numbers(hr_dir, bad_data_hr)

test.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-
from sr.utils.utils import get_config
from sr.logger import logger
from sr.trainer.EEGAN_trainer import EEGAN_Trainer

def main(config):
        
    trainer = EEGAN_Trainer(config,logger)
    trainer.train()
    
    
if __name__ == '__main__':
    config = get_config("config.json")
    main(config)
train.py
zayn309/satellite-images-super-resolution-using-EEGAN-in-pytorch-