text:
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.

"""CSC model."""

import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

import slowfast.models.losses as losses
import slowfast.utils.distributed as du
import slowfast.utils.logging as logging
from slowfast.models.video_model_builder import X3D, MViT, ResNet, SlowFast

from .build import MODEL_REGISTRY

logger = logging.get_logger(__name__)

# Supported model types
_MODEL_TYPES = {
    "slowfast": SlowFast,
    "slow": ResNet,
    "c2d": ResNet,
    "i3d": ResNet,
    "slow_c2d": ResNet,
    "x3d": X3D,
    "mvit": MViT,
}

@MODEL_REGISTRY.register()
class ContrastiveModel(nn.Module):
    """
    Contrastive Model, currently mainly focused on memory bank and CSC.
    """

    def __init__(self, cfg):
        super(ContrastiveModel, self).__init__()
        # Construct the model.
        self.backbone = _MODEL_TYPES[cfg.MODEL.ARCH](cfg)
        self.type = cfg.CONTRASTIVE.TYPE
        self.T = cfg.CONTRASTIVE.T
        self.dim = cfg.CONTRASTIVE.DIM
        self.length = cfg.CONTRASTIVE.LENGTH
        self.k = cfg.CONTRASTIVE.QUEUE_LEN
        self.mmt = cfg.CONTRASTIVE.MOMENTUM
        self.momentum_annealing = cfg.CONTRASTIVE.MOMENTUM_ANNEALING
        self.duration = 1
        self.cfg = cfg
        self.num_gpus = cfg.NUM_GPUS
        self.l2_norm = Normalize()
        self.knn_num_imgs = 0
        self.knn_on = cfg.CONTRASTIVE.KNN_ON
        self.train_labels = np.zeros((0,), dtype=np.int32)
        self.num_pos = 2
        self.num_crops = (
            self.cfg.DATA.TRAIN_CROP_NUM_TEMPORAL
            * self.cfg.DATA.TRAIN_CROP_NUM_SPATIAL
        )
        self.nce_loss_fun = losses.get_loss_func("contrastive_loss")(
            reduction="mean"
        )
        assert self.cfg.MODEL.LOSS_FUNC == "contrastive_loss"
        self.softmax = nn.Softmax(dim=1).cuda()

        if self.type == "mem":
            self.mem_type = cfg.CONTRASTIVE.MEM_TYPE
            if self.mem_type == "1d":
                self.memory = Memory1D(
                    self.length, self.duration, self.dim, cfg
                )
            else:
                self.memory = Memory(self.length, self.duration, self.dim, cfg)
            self.examplar_type = "video"
            self.interp = cfg.CONTRASTIVE.INTERP_MEMORY
        elif self.type == "self":
            pass
        elif self.type == "moco" or self.type == "byol":
            # MoCo components
            self.backbone_hist = _MODEL_TYPES[cfg.MODEL.ARCH](cfg)
            for p in self.backbone_hist.parameters():
                p.requires_grad = False
            self.register_buffer("ptr", torch.tensor([0]))
            self.ptr.requires_grad = False
            stdv = 1.0 / math.sqrt(self.dim / 3)
            self.register_buffer(
                "queue_x",
                torch.rand(self.k, self.dim).mul_(2 * stdv).add_(-stdv),
            )
            self.register_buffer("iter", torch.zeros([1], dtype=torch.long))
            self._batch_shuffle_on = (
                False
                if (
                    "sync" in cfg.BN.NORM_TYPE
                    and cfg.BN.NUM_SYNC_DEVICES == cfg.NUM_GPUS
                )
                or self.type == "byol"
                else True
            )
        elif self.type == "swav":
            self.swav_use_public_code = True
            if self.swav_use_public_code:
                self.swav_prototypes = nn.Linear(
                    self.dim, 1000, bias=False
                )  # for orig implementation
            else:
                self.swav_prototypes = nn.Parameter(
                    torch.randn((self.dim, 1000), dtype=torch.float)
                )
            self.swav_eps_sinkhorn = 0.05
            self.swav_use_the_queue = False
            # optionally starts a queue
            if self.cfg.CONTRASTIVE.SWAV_QEUE_LEN > 0:
                self.register_buffer(
                    "queue_swav",
                    torch.zeros(
                        2,  # = args.crops_for_assign
                        self.cfg.CONTRASTIVE.SWAV_QEUE_LEN
                        // du.get_world_size(),
                        self.dim,
                    ),
                )
        elif self.type == "simclr":
            self._simclr_precompute_pos_neg_mask_multi()
        self.simclr_dist_on = cfg.CONTRASTIVE.SIMCLR_DIST_ON

        # self.knn_mem = Memory1D(self.length, 1, self.dim, cfg) #  does not work
        if self.knn_on:
            self.knn_mem = Memory(self.length, 1, self.dim, cfg)

    @torch.no_grad()
    def knn_mem_update(self, q_knn, index):
        if self.knn_on:
            self.knn_mem.update(
                q_knn,
                momentum=1.0,
                ind=index,
                time=torch.zeros_like(index),
                interp=False,
            )

    @torch.no_grad()
    def init_knn_labels(self, train_loader):
        logger.info("initializing knn labels")
        self.num_imgs = len(train_loader.dataset._labels)
        self.train_labels = np.zeros((self.num_imgs,), dtype=np.int32)
        for i in range(self.num_imgs):
            self.train_labels[i] = train_loader.dataset._labels[i]
        self.train_labels = torch.LongTensor(self.train_labels).cuda()
        if self.length != self.num_imgs:
            logger.error(
                "Kinetics dataloader size: {} differs from memorybank length {}".format(
                    self.num_imgs, self.length
                )
            )
            self.knn_mem.resize(self.num_imgs, 1, self.dim)

    @torch.no_grad()
    def _update_history(self):
        # momentum update
        iter = int(self.iter)
        m = self.mmt
        dist = {}
        for name, p in self.backbone.named_parameters():
            dist[name] = p

        if iter == 0:
            for name, p in self.backbone_hist.named_parameters():
                p.data.copy_(dist[name].data)

        for name, p in self.backbone_hist.named_parameters():
            p.data = dist[name].data * (1.0 - m) + p.data * m

    @torch.no_grad()
    def _batch_shuffle(self, x):
        if len(x) == 2:
            another_crop = True
        else:
            another_crop = False
        if another_crop:
            x, x_crop = x[0], x[1]
        else:
            x = x[0]

        world_size = self.cfg.NUM_GPUS * self.cfg.NUM_SHARDS
        if self.num_gpus > 1:
            if self.cfg.CONTRASTIVE.LOCAL_SHUFFLE_BN:
                x = du.cat_all_gather(x, local=True)
                if another_crop:
                    x_crop = du.cat_all_gather(x_crop, local=True)
                world_size = du.get_local_size()
                gpu_idx = du.get_local_rank()
            else:
                x = du.cat_all_gather(x)
                if another_crop:
                    x_crop = du.cat_all_gather(x_crop)
                gpu_idx = torch.distributed.get_rank()

        idx_randperm = torch.randperm(x.shape[0]).cuda()
        if self.num_gpus > 1:
            torch.distributed.broadcast(idx_randperm, src=0)
        else:
            gpu_idx = 0
        idx_randperm = idx_randperm.view(world_size, -1)
        x = x[idx_randperm[gpu_idx, :]]
        if another_crop:
            x_crop = x_crop[idx_randperm[gpu_idx, :]]

        idx_restore = torch.argsort(idx_randperm.view(-1))
        idx_restore = idx_restore.view(world_size, -1)
        if another_crop:
            return [x, x_crop], idx_restore
        else:
            return [x], idx_restore

    @torch.no_grad()
    def _batch_unshuffle(self, x, idx_restore):
        if self.num_gpus > 1:
            if self.cfg.CONTRASTIVE.LOCAL_SHUFFLE_BN:
                x = du.cat_all_gather(x, local=True)
                gpu_idx = du.get_local_rank()
            else:
                x = du.cat_all_gather(x)
                gpu_idx = torch.distributed.get_rank()
        else:
            gpu_idx = 0

        idx = idx_restore[gpu_idx, :]
        x = x[idx]
        return x

    @torch.no_grad()
    def eval_knn(self, q_knn, knn_k=200):
        with torch.no_grad():
            dist = torch.einsum(
                "nc,mc->nm",
                q_knn.view(q_knn.size(0), -1),
                self.knn_mem.memory.view(self.knn_mem.memory.size(0), -1),
            )
            yd, yi = dist.topk(knn_k, dim=1, largest=True, sorted=True)
        return yd, yi

    def sim_loss(self, q, k):
        similarity = torch.einsum("nc,nc->n", [q, k])  # N-dim
        # similarity += delta_t # higher if time distance is larger
        # sim = sim - max_margin + delta_t * k
        similarity /= self.T  # history-compatible
        loss = -similarity.mean()
        return loss

    @torch.no_grad()
    def momentum_anneal_cosine(self, epoch_exact):
        self.mmt = (
            1
            - (1 - self.cfg.CONTRASTIVE.MOMENTUM)
            * (
                math.cos(math.pi * epoch_exact / self.cfg.SOLVER.MAX_EPOCH)
                + 1.0
            )
            * 0.5
        )

    @torch.no_grad()
    def _dequeue_and_enqueue(self, keys, extra_keys=None):
        ptr = int(self.ptr.item())
        if (
            not self.cfg.CONTRASTIVE.MOCO_MULTI_VIEW_QUEUE
        ):  # TODO: add multiview negatives
            keys_queue_update = [keys[0]]
        else:
            assert (
                len(keys) > 0
            ), "need to have multiple views for adding them to queue"
            keys_queue_update = []
            keys_queue_update += keys
            if extra_keys:
                keys_queue_update += [
                    item for sublist in extra_keys for item in sublist
                ]
        for key in keys_queue_update:
            # write the current feat into queue, at pointer
            num_items = int(key.size(0))

            assert self.k % num_items == 0
            assert ptr + num_items <= self.k
            self.queue_x[ptr : ptr + num_items, :] = key
            # move pointer
            ptr += num_items
            # reset pointer
            if ptr == self.k:
                ptr = 0
            self.ptr[0] = ptr

    @torch.no_grad()
    def batch_clips(self, clips):
        clips_batched = [None] * len(clips[0])
        for i, clip in enumerate(clips):
            for j, view in enumerate(clip):
                if i == 0:
                    clips_batched[j] = view
                else:
                    clips_batched[j] = torch.cat(
                        [clips_batched[j], view], dim=0
                    )
                del view
        return clips_batched

    @torch.no_grad()
    def compute_key_feat(
        self, clips_k, compute_predictor_keys=False, batched_inference=True
    ):
        assert self.training
        # momentum update key encoder
        self._update_history()
        self.iter += 1
        n_clips = len(clips_k)
        bsz = clips_k[0][0].shape[0]
        if n_clips * bsz * clips_k[0][0].numel() > 4 * 64 * 3 * 8 * 224 * 224:
            batched_inference = False  # hack to avoid oom on large inputs
        assert n_clips > 0
        if batched_inference and all(
            [
                clips_k[i][j].shape[1:] == clips_k[0][j].shape[1:]
                for i in range(len(clips_k))
                for j in range(len(clips_k[i]))
            ]
        ):
            clips_k = [self.batch_clips(clips_k)]
            batched = True
        else:
            batched = False

        keys, pred_keys = [], []
        for k in range(0, len(clips_k)):
            clip_k = clips_k[k]
            if self._batch_shuffle_on:
                with torch.no_grad():
                    clip_k, idx_restore = self._batch_shuffle(clip_k)
            with torch.no_grad():
                hist_feat = self.backbone_hist(clip_k)
                if isinstance(hist_feat, list):
                    hist_time = hist_feat[1:]
                    hist_feat = hist_feat[0]
                    if compute_predictor_keys:
                        tks = []
                        for tk in hist_time:
                            tk = self.l2_norm(tk)
                            if self._batch_shuffle_on:
                                tk = self._batch_unshuffle(
                                    tk, idx_restore
                                ).detach()
                            tks.append(tk)
                        pred_keys.append(tks)
                x_hist = self.l2_norm(hist_feat)
                if self._batch_shuffle_on:
                    x_hist = self._batch_unshuffle(x_hist, idx_restore).detach()
            keys.append(x_hist)
        if batched:
            assert len(keys) == 1, "batched input uses single clip"
            batched_key = keys[0]
            if compute_predictor_keys:
                batched_pred_key = pred_keys[0]
            keys, pred_keys = [], []
            for k in range(0, n_clips):
                keys.append(batched_key[k * bsz : (k + 1) * bsz])
                if compute_predictor_keys:
                    pred_keys.append(batched_pred_key[k * bsz : (k + 1) * bsz])
        if compute_predictor_keys:
            return keys, pred_keys
        else:
            return keys

    def forward(
        self, clips, index=None, time=None, epoch_exact=None, keys=None
    ):
        if epoch_exact is not None and self.momentum_annealing:
            self.momentum_anneal_cosine(epoch_exact)

        if self.type == "mem":
            batch_size = clips[0].size(0)
            q = self.backbone(clips)
            if index is None:
                return q
            q = self.l2_norm(q)

            if not self.training:
                assert self.knn_mem.duration == 1
                return self.eval_knn(q)
            time *= self.duration - 1
            clip_ind = torch.randint(
                0,
                self.length,
                size=(
                    batch_size,
                    self.k + 1,
                ),
            ).cuda()
            clip_ind.select(1, 0).copy_(index.data)

            if self.mem_type == "2d":
                if self.interp:
                    time_ind = (
                        torch.empty(batch_size, self.k + 1)
                        .uniform_(0, self.duration - 1)
                        .cuda()
                    )
                else:
                    time_ind = torch.randint(
                        0,
                        self.duration - 1,
                        size=(
                            batch_size,
                            self.k + 1,
                        ),
                    ).cuda()
            else:
                time_ind = torch.zeros(
                    size=(batch_size, self.k + 1), dtype=int
                ).cuda()

            if self.examplar_type == "clip":
                # Diff clip from same video are negative.
                time_ind.select(1, 0).copy_(time.data)
            elif self.examplar_type == "video":
                # Diff clip from same video are positive.
                pass
            else:
                raise NotImplementedError(
                    "unsupported examplar_type {}".format(self.examplar_type)
                )
            k = self.memory.get(clip_ind, time_ind, self.interp)
            # q: N x C, k: N x K x C
            prod = torch.einsum("nc,nkc->nk", q, k)
            prod = torch.div(prod, self.T)

            loss = self.nce_loss_fun(prod)

            self.memory.update(
                q, momentum=self.mmt, ind=index, time=time, interp=self.interp
            )
            self.knn_mem_update(q, index)
            return prod, 0.0, True
        elif self.type == "moco":
            if isinstance(clips[0], list):
                n_clips = len(clips)
                ind_clips = np.arange(
                    n_clips
                )  # clips come ordered temporally from decoder

                clip_q = clips[ind_clips[0]]
                clips_k = [clips[i] for i in ind_clips[1:]]
                # rearange time
                time_q = time[:, ind_clips[0], :]
                time_k = (
                    time[:, ind_clips[1:], :]
                    if keys is None
                    else time[:, ind_clips[0] + 1 :, :]
                )
            else:
                clip_q = clips

            feat_q = self.backbone(clip_q)
            extra_projs = []
            if isinstance(feat_q, list):
                extra_projs = feat_q[1:]
                feat_q = feat_q[0]
                extra_projs = [self.l2_norm(feat) for feat in extra_projs]

            if index is None:
                return feat_q
            q = self.l2_norm(feat_q)
            q_knn = q

            if not self.training:
                return self.eval_knn(q_knn)

            if keys is None:
                keys = self.compute_key_feat(
                    clips_k, compute_predictor_keys=False
                )
                auto_enqueue_keys = True
            else:
                auto_enqueue_keys = False

            # score computation
            queue_neg = torch.einsum(
                "nc,kc->nk", [q, self.queue_x.clone().detach()]
            )

            for k, key in enumerate(keys):
                out_pos = torch.einsum("nc,nc->n", [q, key]).unsqueeze(-1)
                lgt_k = torch.cat([out_pos, queue_neg], dim=1)
                if k == 0:
                    logits = lgt_k
                else:
                    logits = torch.cat([logits, lgt_k], dim=0)

            logits = torch.div(logits, self.T)

            loss = self.nce_loss_fun(logits)
            # update queue
            if self.training and auto_enqueue_keys:
                self._dequeue_and_enqueue(keys)

            self.knn_mem_update(q_knn, index)
            return logits, loss

        elif self.type == "byol":
            clips_key = [None] * len(clips)
            for i, clip in enumerate(clips):
                p = []
                for path in clip:
                    p.append(path)
                clips_key[i] = p
            batch_clips = False
            if isinstance(clips[0], list):
                n_clips = len(clips)
                ind_clips = np.arange(
                    n_clips
                )  # clips come ordered temporally from decoder
                if batch_clips and n_clips > 1:
                    clips_batched = self.batch_clips(clips)
                    clips_key = [clips_batched]
                    clip_q = clips_batched
                else:
                    clip_q = clips[0]
            else:
                clip_q = clips

            feat_q = self.backbone(clip_q)
            predictors = []
            if isinstance(feat_q, list):
                predictors = feat_q[1:]
                feat_q = feat_q[0]
                predictors = [self.l2_norm(feat) for feat in predictors]
            else:
                raise NotImplementedError("BYOL: predictor is missing")
            assert len(predictors) == 1
            if index is None:
                return feat_q
            q = self.l2_norm(feat_q)

            q_knn = q  # projector

            if not self.training:
                return self.eval_knn(q_knn)

            ind_clips = np.arange(
                n_clips
            )  # clips come ordered temporally from decoder

            # rest down is for training
            if keys is None:
                keys = self.compute_key_feat(
                    clips_key, compute_predictor_keys=False
                )

            if self.cfg.CONTRASTIVE.SEQUENTIAL:
                loss_reg = self.sim_loss(predictors[0], keys[0])
                for i in range(1, len(keys)):
                    loss_reg += self.sim_loss(predictors[0], keys[i])
                loss_reg /= len(keys)
            else:
                if batch_clips:
                    bs = predictors[0].shape[0] // 2
                    loss_reg = self.sim_loss(
                        predictors[0][:bs, :], keys[0][bs:, :]
                    ) + self.sim_loss(predictors[0][bs:, :], keys[0][:bs, :])
                    q_knn = q_knn[:bs, :]
                    del clips_batched[0]
                else:
                    loss_q1 = self.sim_loss(predictors[0], keys[1])
                    assert len(clips) == 2
                    clip_q2 = clips[1]
                    feat_q2 = self.backbone(clip_q2)
                    predictors2 = feat_q2[1:]
                    # feat_q2 = feat_q2[0] # not used
                    predictors2 = [self.l2_norm(feat) for feat in predictors2]
                    assert len(predictors2) == 1

                    loss_q2 = self.sim_loss(predictors2[0], keys[0])
                    loss_reg = loss_q1 + loss_q2

            # loss_pos = self.sim_loss(q1_proj, q2_proj)
            dummy_logits = torch.cat(
                (
                    9999.0
                    * torch.ones((len(index), 1), dtype=torch.float).cuda(),
                    torch.zeros((len(index), self.k), dtype=torch.float).cuda(),
                ),
                dim=1,
            )

            self.knn_mem_update(q_knn, index)

            return dummy_logits, loss_reg

        elif self.type == "swav":
            if not isinstance(clips[0], list):
                if self.swav_use_public_code:
                    proj_1, _ = self.run_swav_orig_encoder_q(clips)
                else:
                    proj_1, _ = self.run_swav_encoder_q(clips)
                if index is None:
                    return proj_1
                if not self.training:
                    return self.eval_knn(proj_1)
            n_clips = len(clips)
            ind_clips = np.arange(
                n_clips
            )  # clips come ordered temporally from decoder
            clip_q = clips[0]

            if self.swav_use_public_code:
                # uses official code of SwAV from
                # https://github.com/facebookresearch/swav/blob/master/main_swav.py
                with torch.no_grad():
                    m = self.module if hasattr(self, "module") else self
                    w = m.swav_prototypes.weight.data.clone()
                    w = nn.functional.normalize(w, dim=1, p=2)
                    m.swav_prototypes.weight.copy_(w)

                bs = clips[0][0].size(0)
                output, embedding = [], []

                for i, clip_q in enumerate(clips):
                    x = self.run_swav_orig_encoder_q(clip_q)
                    embedding.append(x[0])
                    output.append(x[1])
                q_knn = embedding[0]
                embedding = torch.cat(embedding, dim=0)
                output = torch.cat(output, dim=0)

                loss_swav = 0
                swav_extra_crops = n_clips - 2
                self.swav_crops_for_assign = np.arange(
                    n_clips - swav_extra_crops
                )
                for i, crop_id in enumerate(self.swav_crops_for_assign):
                    with torch.no_grad():
                        out = output[bs * crop_id : bs * (crop_id + 1)]
                        if (
                            self.cfg.CONTRASTIVE.SWAV_QEUE_LEN > 0
                            and epoch_exact >= 15.0
                        ):
                            if self.swav_use_the_queue or not torch.all(
                                self.queue_swav[i, -1, :] == 0
                            ):
                                self.swav_use_the_queue = True
                                out = torch.cat(
                                    (
                                        torch.mm(
                                            self.queue_swav[i],
                                            m.swav_prototypes.weight.t(),
                                        ),
                                        out,
                                    )
                                )
                            self.queue_swav[i, bs:] = self.queue_swav[
                                i, :-bs
                            ].clone()
                            self.queue_swav[i, :bs] = embedding[
                                crop_id * bs : (crop_id + 1) * bs
                            ]
                        q = out / self.swav_eps_sinkhorn
                        q = torch.exp(q).t()
                        q = (
                            self.distributed_sinkhorn(q, 3)[-bs:]
                            if self.cfg.NUM_SHARDS > 1
                            else self.sinkhorn(q.t(), 3)[-bs:]
                        )
                    subloss = 0
                    for v in np.delete(np.arange(n_clips), crop_id):
                        p = self.softmax(output[bs * v : bs * (v + 1)] / self.T)
                        subloss -= torch.mean(
                            torch.sum(q * torch.log(p), dim=1)
                        )
                    loss_swav += subloss / (n_clips - 1)
                loss_swav /= len(self.swav_crops_for_assign)
            else:
                proj_1, out_1 = self.run_swav_encoder_q(clip_q)
                q_knn = proj_1
                if not self.training:
                    return self.eval_knn(q_knn)
                proj_2, out_2 = self.run_swav_encoder_q(clips[1])
                bs = proj_1.shape[0]
                if self.cfg.CONTRASTIVE.SWAV_QEUE_LEN > 0:
                    if epoch_exact >= 15.0 and not torch.all(
                        self.queue_swav[0, -1, :] == 0
                    ):
                        swav_prototypes = F.normalize(
                            self.swav_prototypes, dim=0, p=2
                        ).detach()
                        out_1 = torch.cat(
                            (
                                torch.mm(
                                    self.queue_swav[0].detach(), swav_prototypes
                                ),
                                out_1,
                            )
                        )
                        out_2 = torch.cat(
                            (
                                torch.mm(
                                    self.queue_swav[1].detach(), swav_prototypes
                                ),
                                out_2,
                            )
                        )
                    # fill the queue
                    self.queue_swav[0, bs:] = self.queue_swav[0, :-bs].clone()
                    self.queue_swav[0, :bs] = proj_1.detach()
                    self.queue_swav[1, bs:] = self.queue_swav[1, :-bs].clone()
                    self.queue_swav[1, :bs] = proj_2.detach()

                with torch.no_grad():
                    code_1 = self.get_code(out_1)
                    code_2 = self.get_code(out_2)
                loss12 = self.KLDivLoss(out_1[-bs:], code_2[-bs:].detach())
                loss21 = self.KLDivLoss(out_2[-bs:], code_1[-bs:].detach())
                loss_swav = loss12 + loss21
            self.knn_mem_update(q_knn, index)
            dummy_logits = torch.cat(
                (
                    9999.0
                    * torch.ones((len(index), 1), dtype=torch.float).cuda(),
                    torch.zeros((len(index), self.k), dtype=torch.float).cuda(),
                ),
                dim=1,
            )
            return dummy_logits, loss_swav

        elif self.type == "simclr":
            if isinstance(clips[0], list):
                n_clips = len(clips)
                clip_q = clips[0]
            else:
                clip_q = clips
            feat_q = self.backbone(clip_q)
            q = self.l2_norm(feat_q)
            if index is None:
                return q
            q_knn = q
            if not self.training:
                return self.eval_knn(q_knn)
            q2 = self.backbone(clips[1])
            q2 = self.l2_norm(q2)
            distributed_loss = False
            if distributed_loss and self.cfg.NUM_GPUS > 1:
                out = torch.cat([q, q2], dim=0)
                if self.cfg.CONTRASTIVE.SIMCLR_DIST_ON:
                    out_all = du.cat_all_gather(out)
                else:
                    out_all = out
                similarity = torch.exp(torch.mm(out, out_all.t()) / self.T)
                Z, loss = 0.0, 0.0
                for loss_id in range(len(self.pos_mask)):
                    pos = torch.sum(similarity * self.pos_mask[loss_id], 1)
                    neg = torch.sum(similarity * self.neg_mask, 1)
                    idx = (
                        1 - torch.sum(self.pos_mask[loss_id], 1) > 0
                    ).detach()
                    term_prob = pos / (pos + neg)
                    term_prob[idx] = 1.0
                    term_loss = torch.log(term_prob)
                    Z += torch.sum(~idx).detach()
                    loss -= torch.sum(term_loss)
                loss /= Z
            else:
                cat_across_gpus = True
                if cat_across_gpus and self.cfg.NUM_GPUS > 1:
                    q = du.AllGatherWithGradient.apply(q)
                    q2 = du.AllGatherWithGradient.apply(q2)
                out = torch.cat([q, q2], dim=0)
                # [2*B, 2*B]
                sim_matrix = torch.exp(
                    torch.mm(out, out.t().contiguous()) / self.T
                )
                # SANITY:
                mask = (
                    torch.ones_like(sim_matrix)
                    - torch.eye(out.shape[0], device=sim_matrix.device)
                ).bool()
                # [2*B, 2*B-1]
                sim_matrix = sim_matrix.masked_select(mask).view(
                    out.shape[0], -1
                )
                # compute loss
                pos_sim = torch.exp(torch.sum(q * q2, dim=-1) / self.T)
                # [2*B]
                pos_sim = torch.cat([pos_sim, pos_sim], dim=0)
                loss = (-torch.log(pos_sim / sim_matrix.sum(dim=-1))).mean()
            self.knn_mem_update(q_knn, index)
            dummy_logits = torch.cat(
                (
                    9999.0
                    * torch.ones((len(index), 1), dtype=torch.float).cuda(),
                    torch.zeros((len(index), self.k), dtype=torch.float).cuda(),
                ),
                dim=1,
            )
            return dummy_logits, loss
        else:
            raise NotImplementedError()

    def _simclr_precompute_pos_neg_mask_multi(self):
        # computed once at the beginning of training
        distributed = self.cfg.CONTRASTIVE.SIMCLR_DIST_ON
        if distributed:
            total_images = self.cfg.TRAIN.BATCH_SIZE * self.cfg.NUM_SHARDS
            world_size = du.get_world_size()
            rank = du.get_rank()
        else:
            total_images = self.cfg.TRAIN.BATCH_SIZE
            world_size = du.get_local_size()
            rank = du.get_local_rank()
        local_orig_images = total_images // world_size
        local_crops = local_orig_images * self.num_crops

        pos_temps = []
        for d in np.arange(self.num_crops):
            pos_temp, neg_temp = [], []
            for i in range(world_size):
                if i == rank:
                    pos = np.eye(local_crops, k=d * local_orig_images) + np.eye(
                        local_crops, k=-local_crops + d * local_orig_images
                    )
                    neg = np.ones((local_crops, local_crops))
                else:
                    pos = np.zeros((local_crops, local_crops))
                    neg = np.zeros((local_crops, local_crops))
                pos_temp.append(pos)
                neg_temp.append(neg)
            pos_temps.append(np.hstack(pos_temp))
            neg_temp = np.hstack(neg_temp)

        pos_mask = []
        for i in range(self.num_crops - 1):
            pos_mask.append(torch.from_numpy(pos_temps[1 + i]))
        neg_mask = torch.from_numpy(neg_temp - sum(pos_temps))

        if self.num_gpus:
            for i in range(len(pos_mask)):
                pos_mask[i] = pos_mask[i].cuda(non_blocking=True)
            neg_mask = neg_mask.cuda(non_blocking=True)
        self.pos_mask, self.neg_mask = pos_mask, neg_mask

    def run_swav_encoder_q(self, im):
        proj = self.backbone(im)  # Nx512, Nx128
        proj = F.normalize(proj, dim=1)  # always normalize
        swav_prototypes = F.normalize(self.swav_prototypes, dim=0, p=2)
        out = proj @ swav_prototypes
        return proj, out

    @torch.no_grad()
    def get_code(self, out):
        with torch.no_grad():
            Q = torch.exp(out / self.swav_eps_sinkhorn)  # BxK
            if self.cfg.NUM_SHARDS > 1:
                Q_sink = self.distributed_sinkhorn(Q.t(), 3)  # BxK
            else:
                Q_sink = self.sinkhorn(Q, 3)  # BxK
        return Q_sink

    def run_swav_orig_encoder_q(self, x):
        x = self.backbone(x)  # Nx512, Nx128
        x = nn.functional.normalize(x, dim=1, p=2)
        if self.swav_prototypes is not None:
            return x, self.swav_prototypes(x)
        return x

    @torch.no_grad()
    def sinkhorn(self, Q, iters):
        with torch.no_grad():
            Q = Q.t()  # KxB
            sum_Q = torch.sum(Q)
            Q /= sum_Q

            r = torch.ones(Q.shape[0]).cuda(non_blocking=True) / Q.shape[0]
            c = torch.ones(Q.shape[1]).cuda(non_blocking=True) / Q.shape[1]

            for _ in range(iters):
                Q *= (r / torch.sum(Q, dim=1)).unsqueeze(1)
                Q *= (c / torch.sum(Q, dim=0)).unsqueeze(0)

            Q = Q / torch.sum(Q, dim=0, keepdim=True)
            return Q.t().float()

    def distributed_sinkhorn(self, Q, nmb_iters):
        with torch.no_grad():
            sum_Q = torch.sum(Q)
            du.all_reduce([sum_Q], average=False)
            Q /= sum_Q

            u = torch.zeros(Q.shape[0]).cuda(non_blocking=True)
            r = torch.ones(Q.shape[0]).cuda(non_blocking=True) / Q.shape[0]
            c = torch.ones(Q.shape[1]).cuda(non_blocking=True) / (
                du.get_world_size() * Q.shape[1]
            )

            curr_sum = torch.sum(Q, dim=1)
            du.all_reduce([curr_sum], average=False)

            for _ in range(nmb_iters):
                u = curr_sum
                Q *= (r / u).unsqueeze(1)
                Q *= (c / torch.sum(Q, dim=0)).unsqueeze(0)
                curr_sum = torch.sum(Q, dim=1)
                du.all_reduce([curr_sum], average=False)
            return (Q / torch.sum(Q, dim=0, keepdim=True)).t().float()

    def KLDivLoss(self, out, code):
        softmax = nn.Softmax(dim=1).cuda()
        p = softmax(out / self.T)
        loss = torch.mean(-torch.sum(code * torch.log(p), dim=1))
        return loss

def l2_loss(x, y):
    return 2 - 2 * (x * y).sum(dim=-1)

class Normalize(nn.Module):
    def __init__(self, power=2, dim=1):
        super(Normalize, self).__init__()
        self.dim = dim
        self.power = power

    def forward(self, x):
        norm = (
            x.pow(self.power).sum(self.dim, keepdim=True).pow(1.0 / self.power)
        )
        out = x.div(norm)
        return out

class Memory(nn.Module):
    def __init__(self, length, duration, dim, cfg):
        super(Memory, self).__init__()
        self.length = length
        self.duration = duration
        self.dim = dim
        stdv = 1.0 / math.sqrt(dim / 3)
        self.register_buffer(
            "memory",
            torch.rand(length, duration, dim).mul_(2 * stdv).add_(-stdv),
        )
        self.device = self.memory.device
        self.l2_norm = Normalize(dim=1)
        self.l2_norm2d = Normalize(dim=2)
        self.num_gpus = cfg.NUM_GPUS

    def resize(self, length, duration, dim):
        self.length = length
        self.duration = duration
        self.dim = dim
        stdv = 1.0 / math.sqrt(dim / 3)
        del self.memory
        self.memory = (
            torch.rand(length, duration, dim, device=self.device)
            .mul_(2 * stdv)
            .add_(-stdv)
            .cuda()
        )

    def get(self, ind, time, interp=False):
        batch_size = ind.size(0)
        with torch.no_grad():
            if interp:
                # mem_idx = self.memory[ind.view(-1), :, :]
                t0 = time.floor().long()  # - 1
                t0 = torch.clamp(t0, 0, self.memory.shape[1] - 1)
                t1 = t0 + 1
                t1 = torch.clamp(t1, 0, self.memory.shape[1] - 1)

                mem_t0 = self.memory[ind.view(-1), t0.view(-1), :]
                mem_t1 = self.memory[ind.view(-1), t1.view(-1), :]
                w2 = time.view(-1, 1) / self.duration
                w_t1 = (time - t0).view(-1, 1).float()
                w_t1 = 1 - w_t1  # hack for inverse
                selected_mem = mem_t0 * (1 - w_t1) + mem_t1 * w_t1
            else:
                # logger.info("1dmem get ind shape {} time shape {}".format(ind.shape, time.shape))
                selected_mem = self.memory[
                    ind.view(-1), time.long().view(-1), :
                ]

        out = selected_mem.view(batch_size, -1, self.dim)
        return out

    def update(self, mem, momentum, ind, time, interp=False):
        if self.num_gpus > 1:
            mem, ind, time = du.all_gather([mem, ind, time])
        with torch.no_grad():
            if interp:
                t0 = time.floor().long()  # - 1
                t0 = torch.clamp(t0, 0, self.memory.shape[1] - 1)
                t1 = t0 + 1
                t1 = torch.clamp(t1, 0, self.memory.shape[1] - 1)
                mem_t0 = self.memory[ind.view(-1), t0.view(-1), :]
                mem_t1 = self.memory[ind.view(-1), t1.view(-1), :]
                w2 = time.float().view(-1, 1) / float(self.duration)
                w_t1 = (time - t0).view(-1, 1).float()
                w_t1 = 1 - w_t1  # hack for inverse

                w_t0 = 1 - w_t1
                # mem = mem.squeeze()
                duo_update = False
                if duo_update:
                    update_t0 = (
                        mem * w_t0 + mem_t0 * w_t1
                    ) * momentum + mem_t0 * (1 - momentum)
                    update_t1 = (
                        mem * w_t1 + mem_t1 * w_t0
                    ) * momentum + mem_t1 * (1 - momentum)
                else:
                    update_t0 = mem * w_t0 * momentum + mem_t0 * (1 - momentum)
                    update_t1 = mem * w_t1 * momentum + mem_t1 * (1 - momentum)

                update_t0 = self.l2_norm(update_t0)
                update_t1 = self.l2_norm(update_t1)

                self.memory[ind.view(-1), t0.view(-1), :] = update_t0.squeeze()
                self.memory[ind.view(-1), t1.view(-1), :] = update_t1.squeeze()
            else:
                mem = mem.view(mem.size(0), 1, -1)
                mem_old = self.get(ind, time, interp=interp)
                mem_update = mem * momentum + mem_old * (1 - momentum)
                mem_update = self.l2_norm2d(mem_update)
                # logger.info("1dmem set ind shape {} time shape {}".format(ind.shape, time.shape))

                # my version
                self.memory[
                    ind.view(-1), time.long().view(-1), :
                ] = mem_update.squeeze()
                return

    def forward(self, inputs):
        pass

class Memory1D(nn.Module):
    def __init__(self, length, duration, dim, cfg):
        super(Memory1D, self).__init__()
        assert duration == 1
        self.length = length
        self.duration = duration
        self.dim = dim
        stdv = 1.0 / math.sqrt(dim / 3)
        self.register_buffer(
            "memory", torch.rand(length, dim).mul_(2 * stdv).add_(-stdv)
        )
        self.l2_norm = Normalize(dim=1)
        self.num_gpus = cfg.NUM_GPUS

    @torch.no_grad()
    def get(self, ind, time, interp=False):
        batch_size = ind.size(0)
        if len(ind.shape) == 1:
            return torch.index_select(self.memory, 0, ind.view(-1)).view(
                batch_size, self.dim
            )
        else:
            return torch.index_select(self.memory, 0, ind.view(-1)).view(
                batch_size, -1, self.dim
            )

    @torch.no_grad()
    def update(self, mem, momentum, ind, time, interp=False):
        if self.num_gpus > 1:
            mem, ind, time = du.all_gather([mem, ind, time])
        mem = mem.view(mem.size(0), -1)
        ind, time = ind.long(), time.long()

        mem_old = self.get(ind, time, interp=interp)
        mem_update = mem_old * (1 - momentum) + mem * momentum
        mem_update = self.l2_norm(mem_update)

        self.memory.index_copy_(0, ind, mem_update)
        return

def cancel_swav_gradients(model, cfg, epoch_exact):
    # cancel some gradients in first epoch of SwAV
    if (
        cfg.MODEL.MODEL_NAME == "ContrastiveModel"
        and cfg.CONTRASTIVE.TYPE == "swav"
        and epoch_exact <= 1.0
    ):
        for name, p in model.named_parameters():
            if "swav_prototypes" in name:
                p.grad = None
    return model

max_stars_repo_name: macsz/SlowFast

max_issues_repo_name: macsz/SlowFast

max_forks_repo_name: macsz/SlowFast


text:
import os
import argparse
from skimage import io
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider, Button, RadioButtons
import psviewer.terastitcher as ts
import psviewer.raw as raw

# Hold the state of the current FOV
c_idx = 0
z_idx = 0
y_idx = 0
x_idx = 0

cs = None
zs = None
ys = None
xs = None
fmt = None

zstep_small = None
zstep_large = None
pixel_width_um = None
clim = None
data_path = None
img = None

ax_img = None
fig, ax = plt.subplots(figsize=(10, 10))
plt.subplots_adjust(left=0.25, bottom=0.2)
ax_cmin = plt.axes([0.25, 0.1, 0.65, 0.02])
ax_cmax = plt.axes([0.25, 0.14, 0.65, 0.02])
ax_auto = plt.axes([0.8, 0.025, 0.1, 0.04])
ax_radio = plt.axes([0.025, 0.5, 0.12, 0.12])
slider_cmin = Slider(ax_cmin, 'Min Intensity', 0, 2**12-1, valinit=0, valstep=1, valfmt='%d')
slider_cmax = Slider(ax_cmax, 'Max Intensity', 0, 2**12-1, valinit=1000, valstep=1, valfmt='%d')
auto_btn = Button(ax_auto, 'Auto')
radio_btn = None

def adjust_cmin(event):
    global clim
    cmin = slider_cmin.val
    cmax = clim[1]
    clim = [cmin, cmax]
    if cmin < cmax:
        ax_img.set_clim(vmin=cmin, vmax=cmax)
    else:
        slider_cmin.set_val(cmax-1)
        ax_img.set_clim(vmin=cmax-1, vmax=cmax)

def adjust_cmax(event):
    global clim
    cmin = clim[0]
    cmax = slider_cmax.val
    clim = [cmin, cmax]
    if cmin < cmax:
        ax_img.set_clim(vmin=cmin, vmax=cmax)
    else:
        slider_cmax.set_val(cmin+1)
        ax_img.set_clim(vmin=cmin, vmax=cmin+1)

def auto_adjust_clim(event):
    global clim
    cmin = img.min()
    cmax = img.max()
    if cmin < cmax:
        clim = [cmin, cmax]
        slider_cmin.set_val(cmin)
        slider_cmax.set_val(cmax)

def change_channel(event):
    global c_idx
    new_c_idx = int(event[-1])
    if new_c_idx != c_idx:
        c_idx = new_c_idx
        plot_img()

slider_cmin.on_changed(adjust_cmin)
slider_cmax.on_changed(adjust_cmax)
auto_btn.on_clicked(auto_adjust_clim)

def plot_img():
    global ax_img, img
    c, z, y, x = cs[c_idx], zs[z_idx], ys[y_idx], xs[x_idx]
    img_path = ts.build_path(data_path, c, z, y, x, fmt)
    if fmt in ['tif', 'tiff']:
        img = io.imread(img_path)
    elif fmt == 'raw':
        img = raw.raw_imread(img_path)
    width = pixel_width_um * img.shape[0]
    x_um, y_um, z_um = x/10, y/10, z/10
    extent = [x_um, x_um+width, y_um+width, y_um]
    if ax_img is None:
        ax_img = ax.imshow(img,
                           extent=extent,
                           clim=clim,
                           cmap='gray',
                           interpolation='none',
                           animated=True)
    else:
        ax_img.set_data(img)
        ax_img.set_extent([x_um, x_um+width, y_um+width, y_um])
    ax.set_title(f'{c}: z-slice ({z_idx+1}/{len(zs)}) {z_um} um')
    ax.set_xlabel(f'x-tile {x_idx+1} / {len(xs)}')
    ax.set_ylabel(f'y-tile {y_idx+1} / {len(ys)}')
    fig.canvas.draw()

def keypress(event):
    global c_idx, z_idx, y_idx, x_idx
    if event.key is ',':
        if z_idx >= zstep_small:
            z_idx -= zstep_small
        else:
            z_idx = 0
        plot_img()
    elif event.key is '.':
        if z_idx < len(zs)-zstep_small:
            z_idx += zstep_small
        else:
            z_idx = len(zs)-1
        plot_img()
    elif event.key is 'pageup':
        if z_idx >= zstep_large:
            z_idx -= zstep_large
        else:
            z_idx = 0
        plot_img()
    elif event.key is 'pagedown':
        if z_idx < len(zs)-zstep_large:
            z_idx += zstep_large
        else:
            z_idx = len(zs)-1
        plot_img()
    elif event.key is 'up':
        if y_idx >= 1:
            y_idx -= 1
        plot_img()
    elif event.key is 'down':
        if y_idx < len(ys)-1:
            y_idx += 1
        plot_img()
    elif event.key is 'left':
        if x_idx >= 1:
            x_idx -= 1
        plot_img()
    elif event.key is 'right':
        if x_idx < len(xs)-1:
            x_idx += 1
        plot_img()
    elif event.key in [str(i) for i in range(len(cs))]:
        c_idx = int(event.key)
        radio_btn.set_active(c_idx)
        plot_img()
    elif event.guiEvent.char in [str(i) for i in range(len(cs))]:  # catch numpad too
        c_idx = int(event.guiEvent.char)
        radio_btn.set_active(c_idx)
        plot_img()

def main():
    global cs, zs, ys, xs, fmt, pixel_width_um, clim, data_path, zstep_small, zstep_large, radio_btn
    parser = argparse.ArgumentParser(description='Pre-stitcher viewer')
    parser.add_argument('--data-path', type=str, help='input path to unstitcher terastitcher volume', default='.')
    parser.add_argument('--zstep-large', type=int, help='number of frames to move in z in large step', default=50)
    parser.add_argument('--zstep-small', type=int, help='number of frames to move in z in small step', default=5)
    parser.add_argument('--pixel-width', '-p', type=float, help='image pixel with in micron', default=1)
    parser.add_argument('--cmin', type=float, help='minimum of grayscale intensity display range', default=0)
    parser.add_argument('--cmax', type=float, help='maximum of grayscale intensity display range', default=1000)
    args = parser.parse_args()

    data_path = os.path.abspath(args.data_path)
    clim = [args.cmin, args.cmax]
    pixel_width_um = args.pixel_width
    zstep_small = args.zstep_small
    zstep_large = args.zstep_large

    fmt = ts.get_format(data_path)
    cs, zs, ys, xs = ts.get_czyx(data_path)

    radio_btn = RadioButtons(ax_radio, cs, active=0)
    radio_btn.on_clicked(change_channel)
    fig.canvas.mpl_connect('key_press_event', keypress)
    plot_img()
    plt.show()

    # TODO: Allow changing the zstep size
    # TODO: cache images for faster rendering

if __name__ == '__main__':
    main()

ext: py

max_stars_repo_name: chunglabmit/psviewer

max_issues_repo_name: chunglabmit/psviewer

max_forks_repo_name: chunglabmit/psviewer

reason: import numpy


text:
# Copyright (c) 2020, The InferLO authors. All rights reserved.
# Licensed under the Apache License, Version 2.0 - see LICENSE file.
import numpy as np
from scipy.special import softmax

from inferlo.pairwise.bruteforce import infer_bruteforce, sample_bruteforce
from inferlo.pairwise.pwf_model import PairWiseFiniteModel

def _stochastic_vector(n):
    x = np.random.random(size=n)
    return x / np.sum(x)

def test_infer_1_variable():
    al_size = 10
    probs = _stochastic_vector(al_size)
    model = PairWiseFiniteModel(1, al_size)
    model.set_field(np.log(probs).reshape(1, al_size))

    result = infer_bruteforce(model)

    assert np.allclose(result.log_pf, 0)
    assert np.allclose(result.marg_prob, probs)

def test_infer_ising_2_variables():
    model = PairWiseFiniteModel(2, 2)
    j = 5 * np.random.random()
    model.add_interaction(0, 1, np.array([[j, -j], [-j, j]]))

    result = model.infer(algorithm='bruteforce')

    assert np.allclose(result.log_pf, np.log(4 * np.cosh(j)))
    assert np.allclose(result.marg_prob, 0.5 * np.ones((2, 2)))

def test_infer_isolated():
    gr_size = 1000
    model = PairWiseFiniteModel(gr_size, 2)
    model.set_field(np.array([[0, 1]] * gr_size))
    res = model.infer(algorithm='bruteforce')
    assert np.allclose(res.log_pf,
                       gr_size * np.log(1 + np.exp(1)))
    assert np.allclose(res.marg_prob - softmax([0, 1]), 0)

def test_max_likelihood():
    model = PairWiseFiniteModel(3, 2)
    model.set_field(np.array([[0.4, 0.6], [0.4, 0.6], [0.4, 0.6]]))
    model.add_interaction(0, 1, np.array([[0, 10], [10, 0]]))
    model.add_interaction(1, 2, np.array([[0, 10], [10, 0]]))

    state = model.max_likelihood(algorithm='bruteforce')

    assert np.allclose(state, np.array([1, 0, 1]))

def test_max_likelihood_isolated():
    gr_size = 1000
    model = PairWiseFiniteModel(gr_size, 2)
    model.set_field(np.array([[0, 1]] * gr_size))
    result = model.max_likelihood(algorithm='bruteforce')
    assert np.allclose(result, np.ones(gr_size))

def test_sample_bruteforce():
    gr_size, num_samples = 3, 50
    model = PairWiseFiniteModel(gr_size, 2)
    model.set_field(np.array([[0, 20]] * gr_size))
    samples = sample_bruteforce(model, num_samples=num_samples)
    assert np.allclose(samples, np.ones((num_samples, gr_size)))

max_stars_repo_name: InferLO/inferlo

max_issues_repo_name: InferLO/inferlo

max_forks_repo_name: InferLO/inferlo

reason: import numpy,from scipy


text:
import os
import numpy as np
import tensorflow as tf

from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import backend as K
from tensorflow.keras.models import Model
from tensorflow.keras.utils import Sequence
from tensorflow.keras.layers import Input, Dense, Conv2D, MaxPooling2D, UpSampling2D, Flatten, Dense, Reshape, Lambda, Conv2DTranspose, UpSampling2D, Concatenate

from utils import *

def adib(latent_dim):

    # Encoder

    encoder_input = keras.Input(shape=(128, 128, 1))

    x = layers.Conv2D(filters=8, kernel_size=(3,5), strides=(2,2), activation=None, padding='same')(encoder_input)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters=16, kernel_size=10, strides=(4,2), activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters=24, kernel_size=10, strides=(4,2), activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters=32, kernel_size=10, strides=(4,4), activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Flatten()(x)
    #x = layers.Dropout(0.3)(x)
    #x = layers.Dense(64, activation="relu")(x)

    z = layers.Dense(latent_dim, name="z")(x)

    encoder = keras.Model(encoder_input, z, name="encoder")

    # Decoder

    latent_inputs = keras.Input(shape=(latent_dim,))

    #x = layers.Dropout(0.5)(latent_inputs)
    dec = layers.Dense(units=1*4*64, activation="relu")(latent_inputs)
    #x = layers.Dropout(0.5)(x)
    dec = layers.Reshape(target_shape=(1,4,64))(dec)
    dec = layers.UpSampling2D(size=(4, 4))(dec)
    dec = layers.Conv2D(filters=24, kernel_size=10, strides=1, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(4, 2))(dec)
    dec = layers.Conv2D(filters=16, kernel_size=10, strides=1, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(4, 2))(dec)
    dec = layers.Conv2D(filters=8, kernel_size=10, strides=1, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(2, 2))(dec)
    decoder_outputs = layers.Conv2D(filters=1, kernel_size=(3,5), strides=1, padding='same', activation='relu')(dec)

    decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")

    # Define VAE

    decoder_output = decoder(encoder(encoder_input))
    model = Model(encoder_input, decoder_output)

    # Loss

    vae_loss = keras.losses.mse(keras.layers.Flatten()(encoder_input),keras.layers.Flatten()(decoder_output))
    vae_loss *= 128*128

    model.add_loss(vae_loss)
    model.compile(optimizer='adam')

    return model

def adib_var(latent_dim):

    # Encoder

    encoder_input = keras.Input(shape=(128, 128, 1))

    x = layers.Conv2D(filters=8, kernel_size=(3,5), strides=(2,2), activation=None, padding='same')(encoder_input)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters=16, kernel_size=10, strides=(4,2), activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters=24, kernel_size=10, strides=(4,2), activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters=32, kernel_size=10, strides=(4,4), activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Flatten()(x)
    #x = layers.Dropout(0.3)(x)
    #x = layers.Dense(64, activation="relu")(x)

    z_mean = layers.Dense(latent_dim, name="z_mean")(x)
    z_log_var = layers.Dense(latent_dim, name="z_log_var")(x)

    z = sampling([z_mean, z_log_var])

    encoder = keras.Model(encoder_input, [z_mean, z_log_var, z], name="encoder")

    # Decoder

    latent_inputs = keras.Input(shape=(latent_dim,))

    #x = layers.Dropout(0.5)(latent_inputs)
    dec = layers.Dense(units=1*4*64, activation="relu")(latent_inputs)
    #x = layers.Dropout(0.5)(x)
    dec = layers.Reshape(target_shape=(1,4,64))(dec)
    dec = layers.UpSampling2D(size=(4, 4))(dec)
    dec = layers.Conv2D(filters=24, kernel_size=10, strides=1, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(4, 2))(dec)
    dec = layers.Conv2D(filters=16, kernel_size=10, strides=1, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(4, 2))(dec)
    dec = layers.Conv2D(filters=8, kernel_size=10, strides=1, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(2, 2))(dec)
    decoder_outputs = layers.Conv2D(filters=1, kernel_size=(3,5), strides=1, padding='same', activation="relu")(dec)

    decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")

    # Define VAE

    decoder_output = decoder(encoder(encoder_input)[2])
    model = Model(encoder_input, decoder_output)

    # Loss

    reconstruction_loss = keras.losses.mse(keras.layers.Flatten()(encoder_input),keras.layers.Flatten()(decoder_output))
    reconstruction_loss *= 128*128

    kl_loss = 1 + z_log_var - K.square(z_mean) - K.exp(z_log_var)
    kl_loss = K.sum(kl_loss, axis=1)
    kl_loss = -0.5 * kl_loss
    vae_loss = K.mean(reconstruction_loss + kl_loss)
    model.add_loss(vae_loss)
    model.compile(optimizer='adam')

    return model

def adib_cond(latent_dim, num_classes):

    # Encoder

    encoder_input = keras.Input(shape=(128, 128, 1))
    classes = Input(shape=(num_classes,))

    x = layers.Conv2D(filters=8, kernel_size=(3,5), strides=(2,2), activation=None, padding='same')(encoder_input)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters=16, kernel_size=9, strides=(4,2), activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters=32, kernel_size=9, strides=(4,2), activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters=64, kernel_size=9, strides=(4,4), activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Flatten()(x)

    x_cond = Concatenate(axis=-1)([x, classes])
    z = Dense(latent_dim, activation='linear', name="z_layer")(x_cond)

    encoder = Model([encoder_input, classes], [z, classes], name="encoder")

    # Decoder

    latent_inputs = keras.Input(shape=(latent_dim,))

    decoder_input = Concatenate(axis=1)([latent_inputs, classes])

    dec = layers.Dense(units=1*4*64, activation="relu")(decoder_input)
    dec = layers.Reshape(target_shape=(1,4,64))(dec)
    dec = layers.Conv2DTranspose(filters=32, kernel_size=9, strides=(4,4), padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.Conv2DTranspose(filters=16, kernel_size=9, strides=(4,2), padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.Conv2DTranspose(filters=8, kernel_size=9, strides=(4,2), padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    decoder_outputs = layers.Conv2DTranspose(filters=1, kernel_size=(3,5), strides=(2,2), padding='same', activation='relu')(dec)

    decoder =  Model([latent_inputs, classes], decoder_outputs, name="decoder")

    # Define VAE

    decoder_output = decoder(encoder([[encoder_input, classes]]))

    model = Model([encoder_input, classes], decoder_output)

    # Loss

    cae_loss = keras.losses.mse(keras.layers.Flatten()(encoder_input),keras.layers.Flatten()(decoder_output))
    cae_loss *= 128*128

    model.add_loss(cae_loss)
    model.compile(optimizer='adam')

    return model

def adib_var_cond(latent_dim, num_classes):

    # Encoder

    encoder_input = keras.Input(shape=(128, 128, 1), name="input")
    classes = keras.Input(shape=(num_classes,), name="class")

    x = layers.Conv2D(filters=8, kernel_size=(3,5), strides=(2,2), activation=None, padding='same')(encoder_input)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters=16, kernel_size=10, strides=(4,2), activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters=24, kernel_size=10, strides=(4,2), activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters=32, kernel_size=10, strides=(4,4), activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Flatten()(x)

    x_cond = Concatenate(axis=-1)([x, classes])

    z_mean = layers.Dense(latent_dim, name="z_mean")(x_cond)
    z_log_var = layers.Dense(latent_dim, name="z_log_var")(x_cond)

    z = sampling([z_mean, z_log_var])

    encoder = Model([encoder_input, classes], [z_mean, z_log_var, z, classes], name="encoder")

    # Decoder

    latent_inputs = keras.Input(shape=(latent_dim,))

    decoder_input = Concatenate(axis=1)([latent_inputs, classes])

    dec = layers.Dense(units=1*4*64, activation="relu")(decoder_input)
    dec = layers.Reshape(target_shape=(1,4,64))(dec)
    dec = layers.UpSampling2D(size=(4, 4))(dec)
    dec = layers.Conv2D(filters=24, kernel_size=10, strides=1, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(4, 2))(dec)
    dec = layers.Conv2D(filters=16, kernel_size=10, strides=1, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(4, 2))(dec)
    dec = layers.Conv2D(filters=8, kernel_size=10, strides=1, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(2, 2))(dec)
    decoder_output = layers.Conv2D(filters=1, kernel_size=(3,5), strides=1, padding='same', activation="relu")(dec)

    decoder =  Model([latent_inputs, classes], decoder_output, name="decoder")

    #pdb.set_trace()

    # Define VAE

    decoder_output = decoder([[encoder([[encoder_input, classes]])[2], classes]])
    model = Model([encoder_input, classes], decoder_output)

    # Loss

    reconstruction_loss = keras.losses.mse(keras.layers.Flatten()(encoder_input),keras.layers.Flatten()(decoder_output))
    reconstruction_loss *= 128*128

    kl_loss = 1 + z_log_var - K.square(z_mean) - K.exp(z_log_var)
    kl_loss = K.sum(kl_loss, axis=1)
    kl_loss = -0.5 * kl_loss
    vae_loss = K.mean(reconstruction_loss + kl_loss)
    model.add_loss(vae_loss)
    model.compile(optimizer='adam')

    return model

def adib_mp(latent_dim):

    # Encoder

    encoder_input = keras.Input(shape=(128, 128, 1))

    x = layers.Conv2D(filters=8, kernel_size=(3,5), activation=None, padding='same')(encoder_input)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(2,2))(x)
    x = layers.Conv2D(filters=16, kernel_size=10, activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(4,2))(x)
    x = layers.Conv2D(filters=24, kernel_size=10, activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(4,2))(x)
    x = layers.Conv2D(filters=32, kernel_size=10, activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(4,2))(x)
    x = layers.Flatten()(x)
    #x = layers.Dropout(0.3)(x)
    #x = layers.Dense(64, activation="relu")(x)

    z = layers.Dense(latent_dim, name="z")(x)

    encoder = keras.Model(encoder_input, z, name="encoder")

    # Decoder

    latent_inputs = keras.Input(shape=(latent_dim,))

    #x = layers.Dropout(0.5)(latent_inputs)
    dec = layers.Dense(units=1*4*64, activation="relu")(latent_inputs)
    #x = layers.Dropout(0.5)(x)
    dec = layers.Reshape(target_shape=(1,4,64))(dec)
    dec = layers.UpSampling2D(size=(4, 4))(dec)
    dec = layers.Conv2D(filters=24, kernel_size=10, strides=1, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(4, 2))(dec)
    dec = layers.Conv2D(filters=16, kernel_size=10, strides=1, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(4, 2))(dec)
    dec = layers.Conv2D(filters=8, kernel_size=10, strides=1, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(2, 2))(dec)
    decoder_outputs = layers.Conv2D(filters=1, kernel_size=(3,5), strides=1, padding='same', activation='relu')(dec)

    decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")

    # Define VAE

    decoder_output = decoder(encoder(encoder_input))
    model = Model(encoder_input, decoder_output)

    # Loss

    vae_loss = keras.losses.mse(keras.layers.Flatten()(encoder_input),keras.layers.Flatten()(decoder_output))
    vae_loss *= 128*128

    model.add_loss(vae_loss)
    model.compile(optimizer='adam')

    return model

def adib_tc(latent_dim):

    # Encoder

    encoder_input = keras.Input(shape=(128, 128, 1))

    x = layers.Conv2D(filters=8, kernel_size=(3,5), strides=(2,2), activation=None, padding='same')(encoder_input)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters=16, kernel_size=10, strides=(4,2), activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters=24, kernel_size=10, strides=(4,2), activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters=32, kernel_size=10, strides=(4,4), activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Flatten()(x)
    #x = layers.Dropout(0.3)(x)
    #x = layers.Dense(64, activation="relu")(x)

    z = layers.Dense(latent_dim, name="z")(x)

    encoder = keras.Model(encoder_input, z, name="encoder")

    # Decoder

    latent_inputs = keras.Input(shape=(latent_dim,))

    #x = layers.Dropout(0.5)(latent_inputs)
    dec = layers.Dense(units=1*4*64, activation="relu")(latent_inputs)
    #x = layers.Dropout(0.5)(x)
    dec = layers.Reshape(target_shape=(1,4,64))(dec)
    dec = layers.Conv2DTranspose(filters=32, kernel_size=10, strides=(4,4), padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.Conv2DTranspose(filters=16, kernel_size=10, strides=(4,2), padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.Conv2DTranspose(filters=8, kernel_size=10, strides=(4,2), padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    decoder_outputs = layers.Conv2DTranspose(filters=1, kernel_size=(3,5), strides=(2,2), padding='same', activation='relu')(dec)

    decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")

    # Define VAE

    decoder_output = decoder(encoder(encoder_input))
    model = Model(encoder_input, decoder_output)

    # Loss

    vae_loss = keras.losses.mse(keras.layers.Flatten()(encoder_input),keras.layers.Flatten()(decoder_output))
    vae_loss *= 128*128

    model.add_loss(vae_loss)
    model.compile(optimizer='adam')

    return model

def adib_mp_tc(latent_dim):

    # Encoder

    encoder_input = keras.Input(shape=(128, 128, 1))

    x = layers.Conv2D(filters=8, kernel_size=(3,5), activation=None, padding='same')(encoder_input)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(2,2))(x)
    x = layers.Conv2D(filters=16, kernel_size=10, activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(4,2))(x)
    x = layers.Conv2D(filters=24, kernel_size=10, activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(4,2))(x)
    x = layers.Conv2D(filters=32, kernel_size=10, activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(4,2))(x)
    x = layers.Flatten()(x)
    #x = layers.Dropout(0.3)(x)
    #x = layers.Dense(64, activation="relu")(x)

    z = layers.Dense(latent_dim, name="z")(x)

    encoder = keras.Model(encoder_input, z, name="encoder")

    # Decoder

    latent_inputs = keras.Input(shape=(latent_dim,))

    #x = layers.Dropout(0.5)(latent_inputs)
    dec = layers.Dense(units=1*4*64, activation="relu")(latent_inputs)
    #x = layers.Dropout(0.5)(x)
    dec = layers.Reshape(target_shape=(1,4,64))(dec)
    dec = layers.Conv2DTranspose(filters=32, kernel_size=10, strides=(4,4), padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.Conv2DTranspose(filters=16, kernel_size=10, strides=(4,2), padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.Conv2DTranspose(filters=8, kernel_size=10, strides=(4,2), padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    decoder_outputs = layers.Conv2DTranspose(filters=1, kernel_size=(3,5), strides=(2,2), padding='same', activation='relu')(dec)

    decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")

    # Define VAE

    decoder_output = decoder(encoder(encoder_input))
    model = Model(encoder_input, decoder_output)

    # Loss

    vae_loss = keras.losses.mse(keras.layers.Flatten()(encoder_input),keras.layers.Flatten()(decoder_output))
    vae_loss *= 128*128

    model.add_loss(vae_loss)
    model.compile(optimizer='adam')

    return model

def adib_filt_2(latent_dim):

    # Encoder

    encoder_input = keras.Input(shape=(128, 128, 1))

    x = layers.Conv2D(filters=8, kernel_size=(3,5), activation=None, padding='same')(encoder_input)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(2,2))(x)
    x = layers.Conv2D(filters=16, kernel_size=9, activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(4,2))(x)
    x = layers.Conv2D(filters=32, kernel_size=9, activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(4,2))(x)
    x = layers.Conv2D(filters=64, kernel_size=9, activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(4,4))(x)
    x = layers.Flatten()(x)
    #x = layers.Dropout(0.3)(x)
    #x = layers.Dense(64, activation="relu")(x)

    z = layers.Dense(latent_dim, name="z")(x)

    encoder = keras.Model(encoder_input, z, name="encoder")

    # Decoder

    latent_inputs = keras.Input(shape=(latent_dim,))

    #x = layers.Dropout(0.5)(latent_inputs)
    dec = layers.Dense(units=1*4*64, activation="relu")(latent_inputs)
    #x = layers.Dropout(0.5)(x)
    dec = layers.Reshape(target_shape=(1,4,64))(dec)
    dec = layers.UpSampling2D(size=(4,4))(dec)
    dec = layers.Conv2D(filters=32, kernel_size=9, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(4,2))(dec)
    dec = layers.Conv2D(filters=16, kernel_size=9, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(4,2))(dec)
    dec = layers.Conv2D(filters=8, kernel_size=9, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(2,2))(dec)
    decoder_outputs = layers.Conv2D(filters=1, kernel_size=(3,5), padding='same', activation='relu')(dec)

    decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")

    # Define VAE

    decoder_output = decoder(encoder(encoder_input))
    model = Model(encoder_input, decoder_output)

    # Loss

    vae_loss = keras.losses.mse(keras.layers.Flatten()(encoder_input),keras.layers.Flatten()(decoder_output))
    vae_loss *= 128*128

    model.add_loss(vae_loss)
    model.compile(optimizer='adam')

    return model

def adib_timbre(latent_dim):

    # Encoder

    encoder_input = keras.Input(shape=(128, 128, 1))

    x_1 = layers.Conv2D(filters=8, kernel_size=(65,3), activation=None, padding='same')(encoder_input)
    x_2 = layers.Conv2D(filters=8, kernel_size=(43,5), activation=None, padding='same')(encoder_input)
    x_3 = layers.Conv2D(filters=8, kernel_size=(21,7), activation=None, padding='same')(encoder_input)
    x_4 = layers.Conv2D(filters=8, kernel_size=(7,21), activation=None, padding='same')(encoder_input)
    x_5 = layers.Conv2D(filters=8, kernel_size=(5,43), activation=None, padding='same')(encoder_input)
    x_6 = layers.Conv2D(filters=8, kernel_size=(3,65), activation=None, padding='same')(encoder_input)

    x_tim = Concatenate(axis=-1)([x_1,x_2,x_3,x_4,x_5,x_6])
    x_tim = layers.Conv2D(filters=8, kernel_size=(1,1), activation=None, padding='same')(x_tim)

    x = layers.BatchNormalization()(x_tim)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(2,2))(x)
    x = layers.Conv2D(filters=16, kernel_size=9, activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(4,2))(x)
    x = layers.Conv2D(filters=32, kernel_size=9, activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(4,2))(x)
    x = layers.Conv2D(filters=64, kernel_size=9, activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(4,4))(x)
    x = layers.Flatten()(x)
    #x = layers.Dropout(0.3)(x)
    #x = layers.Dense(64, activation="relu")(x)

    z = layers.Dense(latent_dim, name="z")(x)

    encoder = keras.Model(encoder_input, z, name="encoder")

    # Decoder

    latent_inputs = keras.Input(shape=(latent_dim,))

    #x = layers.Dropout(0.5)(latent_inputs)
    dec = layers.Dense(units=1*4*64, activation="relu")(latent_inputs)
    #x = layers.Dropout(0.5)(x)
    dec = layers.Reshape(target_shape=(1,4,64))(dec)
    dec = layers.UpSampling2D(size=(4,4))(dec)
    dec = layers.Conv2D(filters=32, kernel_size=9, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(4,2))(dec)
    dec = layers.Conv2D(filters=16, kernel_size=9, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(4,2))(dec)
    dec = layers.Conv2D(filters=8, kernel_size=9, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(2,2))(dec)
    decoder_outputs = layers.Conv2D(filters=1, kernel_size=(3,5), padding='same', activation='relu')(dec)

    decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")

    # Define VAE

    decoder_output = decoder(encoder(encoder_input))
    model = Model(encoder_input, decoder_output)

    # Loss

    vae_loss = keras.losses.mse(keras.layers.Flatten()(encoder_input),keras.layers.Flatten()(decoder_output))
    vae_loss *= 128*128

    model.add_loss(vae_loss)
    model.compile(optimizer='adam')

    return model

def adib_best_cond(latent_dim, num_classes):

    # Encoder

    encoder_input = keras.Input(shape=(128, 128, 1))
    classes = Input(shape=(num_classes,))

    x = layers.Conv2D(filters=8, kernel_size=(3,5), activation=None, padding='same')(encoder_input)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(2,2))(x)
    x = layers.Conv2D(filters=16, kernel_size=9, activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(4,2))(x)
    x = layers.Conv2D(filters=32, kernel_size=9, activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(4,2))(x)
    x = layers.Conv2D(filters=64, kernel_size=9, activation=None, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPool2D(pool_size=(4,4))(x)
    x = layers.Flatten()(x)

    x_cond = Concatenate(axis=-1)([x, classes])
    z = Dense(latent_dim, activation='linear', name="z_layer")(x_cond)

    encoder = Model([encoder_input, classes], [z, classes], name="encoder")

    # Decoder

    latent_inputs = keras.Input(shape=(latent_dim,))

    decoder_input = Concatenate(axis=1)([latent_inputs, classes])

    dec = layers.Dense(units=1*4*64, activation="relu")(decoder_input)
    #x = layers.Dropout(0.5)(x)
    dec = layers.Reshape(target_shape=(1,4,64))(dec)
    dec = layers.UpSampling2D(size=(4,4))(dec)
    dec = layers.Conv2D(filters=32, kernel_size=9, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(4,2))(dec)
    dec = layers.Conv2D(filters=16, kernel_size=9, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(4,2))(dec)
    dec = layers.Conv2D(filters=8, kernel_size=9, padding='same', activation=None)(dec)
    dec = layers.BatchNormalization()(dec)
    dec = layers.ReLU()(dec)
    dec = layers.UpSampling2D(size=(2,2))(dec)
    decoder_outputs = layers.Conv2D(filters=1, kernel_size=(3,5), padding='same', activation='relu')(dec)

    decoder =  Model([latent_inputs, classes], decoder_outputs, name="decoder")

    # Define VAE

    decoder_output = decoder(encoder([[encoder_input, classes]]))

    model = Model([encoder_input, classes], decoder_output)

    # Loss

    cae_loss = keras.losses.mse(keras.layers.Flatten()(encoder_input),keras.layers.Flatten()(decoder_output))
    cae_loss *= 128*128

    model.add_loss(cae_loss)
    model.compile(optimizer='adam')

    return model


text:
#!/usr/bin/env python
# coding: utf-8
# Extended from: https://github.com/mlmed/torchxrayvision/blob/master/scripts/train_model.py
import os,sys
import os,sys,inspect
from glob import glob
from os.path import exists, join
#import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm
import argparse

import torch
import torchvision, torchvision.transforms
import sklearn, sklearn.model_selection

import random
import train_utils
import torchxrayvision as xrv

from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
from padchest_dataset import PC_Dataset_Custom
import multiprocessing

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('-f', type=str, default="", help='')
    parser.add_argument('-name', type=str, default='densenet')
    parser.add_argument('--output_dir', type=str, default=".\outputs")
    parser.add_argument('--dataset', type=str, default="pc")
    parser.add_argument('--dataset_dir', type=str, default="F:\publicimagingdatasets\padchest")
    parser.add_argument('--custom_padchest_csv_file', type=str, default=r"PADCHEST_chest_x_ray_images_labels_160K_01.02.19.csv")
    parser.add_argument('--model', type=str, default="densenet")
    parser.add_argument('--seed', type=int, default=0, help='')
    parser.add_argument('--cuda', type=bool, default=True, help='')
    parser.add_argument('--num_epochs', type=int, default=20, help='')
    parser.add_argument('--batch_size_per_gpu', type=int, default=24, help='')
    parser.add_argument('--shuffle', type=bool, default=True, help='')
    parser.add_argument('--lr', type=float, default=0.001, help='')
    parser.add_argument('--threads', type=int, default=multiprocessing.cpu_count(), help='')
    parser.add_argument('--taskweights', type=bool, default=True, help='')
    parser.add_argument('--featurereg', type=bool, default=False, help='')
    parser.add_argument('--weightreg', type=bool, default=False, help='')
    parser.add_argument('--data_aug', type=bool, default=True, help='')
    parser.add_argument('--data_aug_rot', type=int, default=45, help='')
    parser.add_argument('--data_aug_trans', type=float, default=0.15, help='')
    parser.add_argument('--data_aug_scale', type=float, default=0.15, help='')
    parser.add_argument('--label_concat', type=bool, default=False, help='')
    parser.add_argument('--label_concat_reg', type=bool, default=False, help='')
    parser.add_argument('--multicuda', type=bool, default=True, help='')
    parser.add_argument('--azure_ml', type=bool, default=False, help='')
    parser.add_argument('--testing', type=bool, default=False, help='')

    cfg = parser.parse_args()
    print(cfg)

    #Sets output directory as required by AzureML studio
    if cfg.azure_ml:
        cfg.output_dir = './outputs'

    if cfg.testing:
        cfg.limit = 10
    else:
        cfg.limit = None

    if torch.cuda.device_count() > 1:
        cfg.batch_size_per_gpu =  torch.cuda.device_count() * cfg.batch_size_per_gpu

    data_aug = None
    if cfg.data_aug:
        data_aug = torchvision.transforms.Compose([
            xrv.datasets.ToPILImage(),
            torchvision.transforms.RandomAffine(cfg.data_aug_rot,
                                                translate=(cfg.data_aug_trans, cfg.data_aug_trans),
                                                scale=(1.0-cfg.data_aug_scale, 1.0+cfg.data_aug_scale)),
            torchvision.transforms.ToTensor()
        ])
        print(data_aug)

    # Set resolution according to Densenet
    transforms = torchvision.transforms.Compose([xrv.datasets.XRayCenterCrop(),xrv.datasets.XRayResizer(224)])

    cfg.custom_padchest_csv_file = cfg.dataset_dir + os.sep + cfg.custom_padchest_csv_file
    cfg.dataset_dir = cfg.dataset_dir + os.sep + 'png'
    datas = []
    datas_names = []

    # Customized dataset
    if "pc" in cfg.dataset:
        #Custom dataset
        dataset = PC_Dataset_Custom(
            imgpath=cfg.dataset_dir,
            csvpath=cfg.custom_padchest_csv_file,
            transform=transforms, data_aug=data_aug, unique_patients=False, views=["PA"], flat_dir=False)
        datas.append(dataset)
        datas_names.append("pc")
    print("datas_names", datas_names)

    dataset = xrv.datasets.Merge_Dataset(datas)

    # Setting the seed
    np.random.seed(cfg.seed)
    random.seed(cfg.seed)
    torch.manual_seed(cfg.seed)
    if cfg.cuda:
        torch.cuda.manual_seed_all(cfg.seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

    # You can select different network architectures to train
    if "densenet" in cfg.model:
        model = xrv.models.DenseNet(num_classes=dataset.labels.shape[1], in_channels=1,
                                    **xrv.models.get_densenet_params(cfg.model))
    elif "resnet101" in cfg.model:
        model = torchvision.models.resnet101(num_classes=dataset.labels.shape[1], pretrained=False)
        #patch for single channel
        model.conv1 = torch.nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)

    elif "resnet50" in cfg.model:
        model = torchvision.models.resnet50(num_classes=dataset.labels.shape[1], pretrained=False)
        #patch for single channel
        model.conv1 = torch.nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)

    elif "shufflenet_v2_x2_0" in cfg.model:
        model = torchvision.models.shufflenet_v2_x2_0(num_classes=dataset.labels.shape[1], pretrained=False)
        #patch for single channel
        model.conv1[0] = torch.nn.Conv2d(1, 24, kernel_size=3, stride=2, padding=1, bias=False)
    elif "squeezenet1_1" in cfg.model:
        model = torchvision.models.squeezenet1_1(num_classes=dataset.labels.shape[1], pretrained=False)
        #patch for single channel
        model.features[0] = torch.nn.Conv2d(1, 64, kernel_size=3, stride=2, padding=1, bias=False)
    else:
        raise Exception("no model")

    train_utils.train(model, dataset, cfg)

max_stars_repo_name: Srinivaskatravath/Sample-Imaging-demo