<reponame>duanzhiihao/mycv
from tqdm import tqdm
import random
import torch
import torchvision as tv
import torchvision.transforms.functional as tvf

from mycv.utils.general import parse_config_str
from mycv.datasets.constants import IMAGENET_MEAN, IMAGENET_STD


def get_tv_interpolation(name='bilinear'):
    """ get interpolation for torchvision

    Args:
        name (str, optional): 'bilinear' or 'bicubic'
    """
    if name == 'bilinear':
        # if float(tv.__version__[:-2]) < 0.10:
        #     interp = Image.BILINEAR
        #     print(f'Warning: torchvision version: {tv.__version__} do not support InterpolationMode.',
        #         'Using PIL.Image int instead...')
        interp = tv.transforms.InterpolationMode.BILINEAR
    elif name == 'bicubic':
        # if float(tv.__version__[:-2]) < 0.10:
        #     interp = Image.BICUBIC
        #     print(f'Warning: torchvision version: {tv.__version__} do not support InterpolationMode.',
        #         'Using PIL.Image int instead...')
        interp = tv.transforms.InterpolationMode.BICUBIC
    else:
        raise NotImplementedError()
    return interp


class IdentityTransform():
    def __call__(self, x):
        return x


class RandomInterpolation():
    value = ('bilinear', 'bicubic')


class RandomResizedCropInterp(tv.transforms.RandomResizedCrop):
    def __init__(self, size, scale=(0.08, 1.0), ratio=(3. / 4., 4. / 3.)):
        super().__init__(size, scale=scale, ratio=ratio, interpolation=None)
        self.interpolation = RandomInterpolation()

    def forward(self, img):
        interp = random.choice(self.interpolation.value)
        interp = get_tv_interpolation(interp)
        i, j, h, w = self.get_params(img, self.scale, self.ratio)
        return tvf.resized_crop(img, i, j, h, w, self.size, interp)


def get_input_normalization(name: str):
    """ get tv.transforms.Normalize according to the name

    Args:
        name (str): 'imagenet', 'inception'
    """
    if name == 'imagenet':
        norm = tv.transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
    elif name == 'inception':
        norm = tv.transforms.Normalize(mean=(0.5,0.5,0.5), std=(0.5,0.5,0.5))
    elif not name:
        norm = IdentityTransform()
    else:
        raise ValueError(f'Unknown input_norm: {name}')
    return norm


def get_transform(tf_str: str, img_size: int, input_norm):
    """ Get transform given a config string

    Args:
        tf_str (str): transform congif string
        img_size (int): input image size
        input_norm ([type]): input normalization type
    """
    aug_dict = parse_config_str(tf_str)

    if aug_dict.pop('finetune', False):
        transform = [
            tv.transforms.RandomResizedCrop(img_size, scale=(0.64, 1)),
            tv.transforms.RandomHorizontalFlip(),
        ]
    else:
        transform = [
            # tv.transforms.RandomResizedCrop(img_size),
            RandomResizedCropInterp(img_size),
            tv.transforms.RandomHorizontalFlip(),
        ]

    if aug_dict.pop('baseline', False):
        transform.append(tv.transforms.ToTensor())

    if aug_dict.pop('rand', False):
        # from timm.data.transforms import RandomResizedCropAndInterpolation
        from timm.data.random_erasing import RandomErasing
        from timm.data.auto_augment import rand_augment_transform

        assert input_norm == 'imagenet', f'input_norm required to be imagenet but get {input_norm}'
        if aug_dict.pop('inc', False):
            ra_str = 'rand-m9-mstd0.5-inc1'
        else:
            ra_str = 'rand-m9-mstd0.5'
        reprob = aug_dict.pop('re', 0.5)
        transform = [
            # RandomResizedCropAndInterpolation(img_size, interpolation='random'),
            RandomResizedCropInterp(img_size),
            tv.transforms.RandomHorizontalFlip(),
            rand_augment_transform(
                config_str=ra_str,
                hparams={'translate_const': int(img_size*0.45),
                         'img_mean': tuple([round(255 * x) for x in IMAGENET_MEAN])}
            ),
            tv.transforms.ToTensor(),
            RandomErasing(reprob, mode='pixel', max_count=1, device='cpu')
        ]
        assert len(aug_dict) == 0

    if aug_dict.pop('trivial', False):
        from torchvision.transforms.autoaugment import TrivialAugmentWide
        transform.extend([
            TrivialAugmentWide(interpolation=get_tv_interpolation('bilinear')),
            tv.transforms.PILToTensor(),
            tv.transforms.ConvertImageDtype(torch.float)
        ])

    if input_norm:
        norm = get_input_normalization(input_norm)
        transform.append(norm)

    reprob = aug_dict.pop('re', 0)
    if reprob > 0:
        transform.append(tv.transforms.RandomErasing(p=reprob))

    if len(aug_dict) != 0:
        raise ValueError(f'Unknown transform augmentation str: {tf_str}')

    transform = tv.transforms.Compose(transform)
    return transform


def get_trainloader(root_dir, aug: str, img_size: int, input_norm: str,
                    batch_size: int, workers: int, distributed=False):
    """ get training data loader

    Args:
        root_dir ([type]): [description]
        aug (str): transform congif string
        img_size (int): input image size
        input_norm ([type]): input normalization type
        batch_size (int): [description]
        workers (int): [description]
    """
    transform = get_transform(aug, img_size, input_norm)

    trainset = tv.datasets.ImageFolder(root=root_dir, transform=transform)
    sampler = torch.utils.data.distributed.DistributedSampler(trainset) if distributed else None
    trainloader = torch.utils.data.DataLoader(
        trainset, batch_size=batch_size, shuffle=(sampler is None), num_workers=workers,
        pin_memory=True, drop_last=False, sampler=sampler
    )
    return trainloader


@torch.no_grad()
def imcls_evaluate(model: torch.nn.Module, testloader):
    """ Image classification evaluation with a testloader.

    Args:
        model (torch.nn.Module): pytorch model
        testloader (torch.utils.data.Dataloader): test dataloader
    """
    device = next(model.parameters()).device
    dtype = next(model.parameters()).dtype
    forward_ = getattr(model, 'forward_cls', model.forward)
    nC = len(testloader.dataset.classes)

    tp1_sum, tp5_sum, total_num = 0, 0, 0
    print(f'Evaluating {type(model)}, device={device}, dtype={dtype}')
    print(f'batch_size={testloader.batch_size}, num_workers={testloader.num_workers}')
    pbar = tqdm(testloader)
    for imgs, labels in pbar:
        # sanity check
        imgs: torch.FloatTensor
        assert (imgs.dim() == 4)
        nB = imgs.shape[0]
        labels: torch.LongTensor
        assert (labels.shape == (nB,)) and (labels.dtype == torch.int64)
        assert 0 <= labels.min() and labels.max() <= nC-1

        # forward pass, get prediction
        # _debug(imgs)
        imgs = imgs.to(device=device, dtype=dtype)
        p = forward_(imgs)
        # reduce to top5
        _, top5_idx = torch.topk(p, k=5, dim=1, largest=True)
        # compute top1 and top5 matches (true positives)
        correct5 = torch.eq(labels.view(nB, 1), top5_idx.cpu())
        _tp1 = correct5[:, 0].sum().item()
        _tp5 = correct5.any(dim=1).sum().item()
        tp1_sum += _tp1
        tp5_sum += _tp5

        # logging
        total_num += nB
        msg = f'top1: {tp1_sum/total_num:.4g}, top5: {tp5_sum/total_num:.4g}'
        pbar.set_description(msg)
    pbar.close()

    # compute total statistics and return
    assert total_num == len(testloader.dataset)
    acc_top1 = tp1_sum / total_num
    acc_top5 = tp5_sum / total_num
    results = {'top1': acc_top1, 'top5': acc_top5}
    return results

def _debug(imgs):    
    import matplotlib.pyplot as plt
    input_mean = torch.FloatTensor(IMAGENET_MEAN).view(3, 1, 1)
    input_std  = torch.FloatTensor(IMAGENET_STD).view(3, 1, 1)
    im = imgs[0] * input_std + input_mean
    im = im.permute(1,2,0).numpy()
    plt.imshow(im); plt.show()


def main():
    transform = RandomResizedCropInterp(224)
    debug = 1


if __name__ == '__main__':
    main()

