Code for ICCV 2023 paper "Contrastive Pseudo Learning for Open-World DeepFake Attribution"
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
import os
import argparse
import lightning.pytorch as pl
from glob import glob
from omegaconf import OmegaConf
from loguru import logger

import owdfa.algorithm as algorithm
from owdfa.datasets import create_dataloader

import warnings
warnings.filterwarnings("ignore")


def main():
    # set configs
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--config', type=str,
                        default='./configs/eval.yaml')
    parser.add_argument('--exam_id', type=str, default='')
    parser.add_argument('--ckpt_path', type=str, default='')
    parser.add_argument('--output_log', type=str, default='eval.log')
    parser.add_argument('--split', type=str, default='test')
    parser.add_argument('--distributed', type=int, default=0)
    parser.add_argument('--debug', action='store_true', default=False)
    args = parser.parse_args()
    local_config = OmegaConf.load(args.config)
    for k, v in local_config.items():
        setattr(args, k, v)

    os.environ['TORCH_HOME'] = args.torch_home

    # search the checkpoint file according EXAM ID
    if args.exam_id:
        ckpt_path = glob(f'wandb/*{args.exam_id}/ckpts/*.ckpt')
        if len(ckpt_path) >= 1:
            ckpt_path = sorted(ckpt_path)
            args.ckpt_path = ckpt_path[-1]
    exam_dir = os.path.dirname(os.path.dirname(args.ckpt_path))

    # add log file
    if len(args.output_log) > 0:
        logger.add(f'{exam_dir}/{args.output_log}', level="INFO")

    # load dataset
    test_dataloader = create_dataloader(args, split=args.split)

    method = algorithm.__dict__[args.method.name](args)

    trainer = pl.Trainer(default_root_dir=exam_dir)
    trainer.validate(method, test_dataloader, args.ckpt_path)


if __name__ == '__main__':
    main()

eval.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
from .sl import SLModel
from .owssl import CPL

__all__ = ['SLModel', 'CPL']

owdfa/algorithm/__init__.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from owdfa.algorithm.ssl import SSLModel
from owdfa.losses import entropy, MarginLoss


class CPL(SSLModel):
    def __init__(self, args):
        super(CPL, self).__init__(args)
        self.ce = MarginLoss(m=-1 * args.method.CPL.m, s=1)
        self.glv = nn.BCELoss()
        self.eta1 = args.method.CPL.eta1
        self.eta2 = args.method.CPL.eta2
        self.eta3 = args.method.CPL.eta3
        self.batch_size = args.train.batch_size

    def get_loss_names(self):
        loss_name = ['total_loss', 'ce_loss',
                     'glv_loss', 'csp_loss', 'en_loss']
        return loss_name

    def _get_pair(self, target, f_g, f_p, labeled_len, total_len):
        pos_pairs = []
        target_np = target.cpu().numpy()
        # label part
        for ind in range(labeled_len):
            target_i = target_np[ind]
            idxs = np.where(target_np == target_i)[0]
            if len(idxs) == 1:
                pos_pairs.append(idxs[0])
            else:
                selec_idx = np.random.choice(idxs, 1)
                while selec_idx == ind:
                    selec_idx = np.random.choice(idxs, 1)
                pos_pairs.append(int(selec_idx))
        # unlabel part
        # 1. global feature
        feat_detach = f_g.detach()
        feat_g_norm = feat_detach / torch.norm(feat_detach, 2, 1, keepdim=True)
        global_cosine_dist = torch.mm(feat_g_norm, feat_g_norm.t())
        global_sim = global_cosine_dist
        # 2. local feature
        part_feat = f_p.detach()
        part_norm = torch.norm(part_feat, 2, 1)
        feat_p_norm = part_feat / torch.norm(part_feat, 2, 1, keepdim=True)
        part_cosine_dist = torch.bmm(feat_p_norm.permute(
            2, 0, 1), feat_p_norm.permute(2, 1, 0)).permute(1, 2, 0)
        # 3. partial mask
        part_norm = part_norm / torch.norm(part_norm, 2, 1, keepdim=True)
        part_sim = (part_cosine_dist * part_norm.repeat(total_len,
                    1, 1).permute(1, 0, 2)).sum(dim=2)

        # Global similarity & Part filter
        _, pos_idx = torch.topk(global_sim[labeled_len:, :], 2, dim=1)
        vals, _ = torch.topk(part_sim[labeled_len:, :], 2, dim=1)
        choose_k = 2  # this parameter should be fine-tuned with different task
        max_pos = torch.topk(part_sim[:, pos_idx[:, 1]], choose_k, dim=0)[
            0][choose_k-1]
        mask_1 = (vals[:, 1] - max_pos).ge(0).float()
        mask_0 = (vals[:, 1] - max_pos).lt(0).float()
        pos_idx_1 = (pos_idx[:, 1] * mask_1).cpu().numpy()
        pos_idx_0 = (pos_idx[:, 0] * mask_0).cpu().numpy()
        pos_idx = (pos_idx_1 + pos_idx_0).flatten().tolist()
        pos_pairs.extend(pos_idx)
        return pos_pairs, mask_1

    def ce_loss(self, output, target, labeled_len):
        return self.ce(output[:labeled_len], target)

    def glv_loss(self, prob, **kwargs):
        target = kwargs['target']
        feat = kwargs['feat']
        feat_p = kwargs['feat_p']
        labeled_len = kwargs['labeled_len']
        total_len = kwargs['total_len']

        pos_pairs, mask_1 = self._get_pair(
            target, feat, feat_p, labeled_len, total_len)
        tar_prob = prob[pos_pairs, :]
        if len(tar_prob) < self.batch_size:
            return torch.zeros(1).cuda(), mask_1
        tar_sim = torch.bmm(prob.view(self.batch_size, 1, -1),
                            tar_prob.view(self.batch_size, -1, 1)).squeeze()
        tar_ones = torch.ones_like(tar_sim)
        return self.glv(tar_sim, tar_ones), mask_1

    def csp_loss(self, output, prob_ul, labeled_len, mask_1):
        rand_prob = F.gumbel_softmax(
            output[labeled_len:].detach(), tau=1, hard=False)
        target_u_pl = torch.argmax(rand_prob, dim=1)
        max_prob_pl = prob_ul.gather(1, target_u_pl.view(-1, 1)).squeeze()
        csp_loss = (F.cross_entropy(
            output[labeled_len:], rand_prob, reduction='none') * max_prob_pl * mask_1).mean()
        return csp_loss

    def en_loss(self, prob):
        return entropy(torch.mean(prob, 0))

    def loss(self, **kwargs):
        labeled_len = kwargs['labeled_len']
        output = kwargs['output']
        target = kwargs['target']

        loss_map = {}
        # Cross Entropy loss
        ce_loss = self.ce_loss(output, target, labeled_len)
        loss_map['ce_loss'] = ce_loss
        # Global Local Voting loss
        prob = F.softmax(output, dim=1)
        glv_loss, mask = self.glv_loss(prob, **kwargs)
        loss_map['glv_loss'] = glv_loss
        # Confident-based Soft Pseudo label
        prob_ul = F.softmax(output[labeled_len:], dim=1)
        csp_loss = self.csp_loss(output, prob_ul, labeled_len, mask)
        loss_map['csp_loss'] = csp_loss
        # Entropy loss
        en_loss = self.en_loss(prob)
        loss_map['en_loss'] = en_loss
        # Total loss
        total_loss = ce_loss + self.eta1 * glv_loss + \
            self.eta2 * csp_loss - self.eta3 * en_loss
        loss_map['total_loss'] = total_loss
        return loss_map

owdfa/algorithm/owssl/CPL.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
from .CPL import CPL

__all__ = ['CPL']

owdfa/algorithm/owssl/__init__.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
from .maxmatch import MaxHook
from .fixmatch import FixMatchHook
from .flexmatch import FlexMatchHook
from .freematch import FreeMatchHook

__all__ = ['MaxHook', 'FixMatchHook', 'FlexMatchHook', 'FreeMatchHook']

owdfa/algorithm/pseudo_hook/__init__.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
import torch
import torch.nn.functional as F


class PseudoHook():
    """PseudoHook

    Hook for pseudo label loss

    Parameters
    ----------
    eta: float
        Weight of pseudo label loss
    """

    def __init__(self, **kwargs):
        raise NotImplementedError

    def masking(self, logits_x_ulb, **kwargs):
        raise NotImplementedError

    def pce_loss(self, **kwargs):
        """Compute pseudo label loss

        Parameters
        ----------
        output: torch.Tensor
            Output of model
        labeled_len: int
            Length of labeled data
        idx_ulb: torch.Tensor
            Index of unlabeled data

        Returns
        -------
        pce_loss: torch.Tensor
            Pseudo label loss
        """
        output = kwargs['output']
        labeled_len = kwargs['labeled_len']
        prob_ul = F.softmax(output[labeled_len:], dim=1)
        mask_pl = self.masking(prob_ul, **kwargs)
        _, targets_u_pl = torch.max(prob_ul.detach(), dim=1)
        pce_loss = (F.cross_entropy(
            output[labeled_len:], targets_u_pl, reduction='none') * mask_pl).mean()
        return pce_loss * self.eta

owdfa/algorithm/pseudo_hook/base.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
import torch

from .base import PseudoHook


class FixMatchHook(PseudoHook):
    """
    Threshold in FixMatch
    """

    def __init__(self, **kwargs):
        self.p_cutoff = kwargs['p_cutoff'] if 'p_cutoff' in kwargs else 0.95
        self.eta = kwargs['eta'] if 'eta' in kwargs else 1.0

    @torch.no_grad()
    def masking(self, logits_x_ulb, **kwargs):
        probs_x_ulb = logits_x_ulb.detach()
        max_probs, _ = torch.max(probs_x_ulb, dim=-1)
        mask = max_probs.ge(self.p_cutoff)
        mask = mask.to(max_probs.dtype)
        return mask

owdfa/algorithm/pseudo_hook/fixmatch.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
from collections import Counter
from copy import deepcopy
import torch

from .base import PseudoHook


class FlexMatchHook(PseudoHook):
    """
    Adaptive Thresholding in FlexMatch
    """

    def __init__(self, **kwargs):
        self.num_classes = kwargs['num_classes']
        self.ulb_dest_len = kwargs['ulb_dest_len']
        self.p_cutoff = kwargs['p_cutoff'] if 'p_cutoff' in kwargs else 0.95
        self.thresh_warmup = kwargs['thresh_warmup'] if 'thresh_warmup' in kwargs else True
        self.eta = kwargs['eta'] if 'eta' in kwargs else 1.0

        self.selected_label = torch.ones(
            (self.ulb_dest_len,), dtype=torch.long, ) * -1
        self.classwise_acc = torch.zeros((self.num_classes,))

    @torch.no_grad()
    def update(self, **kwargs):
        pseudo_counter = Counter(self.selected_label.tolist())

        if max(pseudo_counter.values()) < self.ulb_dest_len:  # not all(5w) -1
            if self.thresh_warmup:
                for i in range(self.num_classes):
                    self.classwise_acc[i] = pseudo_counter[i] / \
                        max(pseudo_counter.values())
            else:
                wo_negative_one = deepcopy(pseudo_counter)
                if -1 in wo_negative_one.keys():
                    wo_negative_one.pop(-1)
                for i in range(self.num_classes):
                    self.classwise_acc[i] = pseudo_counter[i] / \
                        max(wo_negative_one.values())

    @torch.no_grad()
    def masking(self, logits_x_ulb, idx_ulb, **kwargs):
        if not self.selected_label.is_cuda:
            self.selected_label = self.selected_label.to(logits_x_ulb.device)
        if not self.classwise_acc.is_cuda:
            self.classwise_acc = self.classwise_acc.to(logits_x_ulb.device)

        probs_x_ulb = logits_x_ulb.detach()
        max_probs, max_idx = torch.max(probs_x_ulb, dim=-1)
        mask = max_probs.ge(
            self.p_cutoff * (self.classwise_acc[max_idx] / (2. - self.classwise_acc[max_idx])))  # convex
        select = max_probs.ge(self.p_cutoff)
        mask = mask.to(max_probs.dtype)

        # update
        if idx_ulb[select == 1].nelement() != 0:
            self.selected_label[idx_ulb[select == 1]] = max_idx[select == 1]
        self.update()
        return mask

owdfa/algorithm/pseudo_hook/flexmatch.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
import torch

from .base import PseudoHook


class FreeMatchHook(PseudoHook):
    """
    SAT in FreeMatch
    """

    def __init__(self, **kwargs):
        self.num_classes = kwargs['num_classes']
        self.m = kwargs['momentum'] if 'momentum' in kwargs else 0.999
        self.eta = kwargs['eta'] if 'eta' in kwargs else 1.0

        self.p_model = torch.ones((self.num_classes))  # / self.num_classes
        self.label_hist = torch.ones((self.num_classes))  # / self.num_classes
        self.time_p = self.p_model.mean()

    def value(self):
        mod = self.p_model / torch.max(self.p_model, dim=-1)[0]
        return self.time_p * mod

    @torch.no_grad()
    def update(self, **kwargs):
        probs_x_ulb = kwargs['probs_x_ulb']
        max_probs, max_idx = torch.max(probs_x_ulb, dim=-1, keepdim=True)
        self.time_p = self.time_p * self.m + \
            (1 - self.m) * torch.quantile(max_probs, 0.8)
        self.p_model = self.p_model * self.m + \
            (1 - self.m) * probs_x_ulb.mean(dim=0)
        hist = torch.bincount(
            max_idx.reshape(-1), minlength=self.p_model.shape[0]).to(self.p_model.dtype)
        self.label_hist = self.label_hist * self.m + \
            (1 - self.m) * (hist / hist.sum())

    @torch.no_grad()
    def masking(self, logits_x_ulb, **kwargs):
        if not self.p_model.is_cuda:
            self.p_model = self.p_model.to(logits_x_ulb.device)
        if not self.label_hist.is_cuda:
            self.label_hist = self.label_hist.to(logits_x_ulb.device)
        if not self.time_p.is_cuda:
            self.time_p = self.time_p.to(logits_x_ulb.device)

        probs_x_ulb = logits_x_ulb.detach()
        self.update(probs_x_ulb=probs_x_ulb)

        max_probs, max_idx = probs_x_ulb.max(dim=-1)
        mod = self.p_model / torch.max(self.p_model, dim=-1)[0]
        mask = max_probs.ge(self.time_p * mod[max_idx]).to(max_probs.dtype)
        return mask

owdfa/algorithm/pseudo_hook/freematch.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
import torch

from .base import PseudoHook


class MaxHook(PseudoHook):
    """
    Default selection in MaxMatch
    """

    def __init__(self, **kwargs):
        self.eta = kwargs['eta'] if 'eta' in kwargs else 1.0

    @torch.no_grad()
    def masking(self, logits_x_ulb, **kwargs):
        probs_x_ulb = logits_x_ulb.detach()
        max_probs, _ = torch.max(probs_x_ulb, dim=-1)
        return torch.ones_like(max_probs)

owdfa/algorithm/pseudo_hook/maxmatch.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
import torch.nn.functional as F
import lightning.pytorch as pl
import wandb
from loguru import logger
from timm.models import resume_checkpoint

from owdfa.utils import gather_tensor, val_stat
import owdfa.encoder as encoder
import owdfa.optimizers as optimizers
import owdfa.schedulers as schedulers
import owdfa.losses as losses


class SLModel(pl.LightningModule):
    def __init__(self, args, **kwargs):
        super().__init__()
        self.args = args
        self.encoder = encoder.__dict__[args.model.name](**args.model.params)
        if args.model.resume is not None:
            resume_checkpoint(self.encoder, args.model.resume)
            if args.local_rank == 0:
                logger.info(f'resume model from {args.model.resume}')

    def configure_optimizers(self):
        optimizer = optimizers.__dict__[self.args.optimizer.name](
            self.encoder.parameters(), **self.args.optimizer.params)
        scheduler = schedulers.__dict__[self.args.scheduler.name](
            optimizer, **self.args.scheduler.params)
        return [optimizer], [scheduler]

    def training_step(self, batch, batch_idx):
        criterion = losses.__dict__[
            self.args.loss.name](**self.args.loss.params)

        images = batch['image']
        targets = batch['target']

        output = self.encoder(images)
        loss = criterion(output, targets)
        self.log('ce_loss', loss, on_step=True, prog_bar=True)
        return loss

    def on_validation_epoch_start(self):
        self.val_step_outputs = {
            'tags': [],
            'preds': [],
            'label': [],
            'conf': [],
        }

    def validation_step(self, batch, batch_idx):
        tags = batch['tag']
        images = batch['image']
        targets = batch['target']

        image = images[tags != 0]
        target = targets[tags != 0]
        tag = tags[tags != 0]

        output = self.encoder(image)
        prob = F.softmax(output, dim=1)
        conf, pred = prob.max(1)

        self.val_step_outputs['tags'].extend(tag)
        self.val_step_outputs['preds'].extend(pred)
        self.val_step_outputs['label'].extend(target)
        self.val_step_outputs['conf'].extend(conf)
        return pred

    def on_validation_epoch_end(self):
        logger.info(f'Epoch-{self.current_epoch} validation finished')
        y_tags = gather_tensor(self.val_step_outputs['tags'],
                               dist_=self.args.distributed, to_numpy=True).astype(int)
        y_pred = gather_tensor(self.val_step_outputs['preds'],
                               dist_=self.args.distributed, to_numpy=True).astype(int)
        y_label = gather_tensor(self.val_step_outputs['label'],
                                dist_=self.args.distributed, to_numpy=True).astype(int)
        y_conf = gather_tensor(self.val_step_outputs['conf'],
                               dist_=self.args.distributed, to_numpy=True)
        results = val_stat(y_tags, y_pred, y_label, y_conf)

        self.log_dict(results, on_epoch=True)
        try:
            wandb.log(results, step=self.current_epoch)
        except:
            pass

    def predict_step(self, batch, batch_idx):
        images = batch['image']
        feature, _ = self.encoder.forward_features(images)
        feature = feature.detach().cpu().numpy()

        return {
            'feature': feature,
            'tag': batch['tag'],
            'target': batch['target'],
            'img_path': batch['img_path'],
        }

owdfa/algorithm/sl.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
import torch
import torch.nn.functional as F
import wandb

from owdfa.algorithm.sl import SLModel
import owdfa.algorithm.pseudo_hook as pseudo_hook
from owdfa.utils import AverageMeter, update_meter


class SSLModel(SLModel):
    def __init__(self, args):
        super().__init__(args)
        hook = None
        if 'pseudo' in args and args.pseudo.hook != 'None':
            hook = pseudo_hook.__dict__[args.pseudo.hook](ulb_dest_len=len(
                self.train_dataloader.dataset), **args.pseudo.params)
        self.hook = hook

    def get_loss_names(self) -> list:
        """Obtain the name of all losses

        Returns
        -------
        loss_names: list
            List with the name of all losses
        """
        raise NotImplementedError

    def loss(self, **kwargs) -> dict:
        """Compute all losses

        Parameters
        ----------
        labeled_len: int
            Length of labeled data
        output: torch.Tensor
            Output of model
        feat: torch.Tensor
            Feature of model
        target: torch.Tensor
            Ground truth
        device: torch.device
            Device
        kwargs: dict
            Other arguments

        Returns
        -------
        loss: dict
            Dictionary with all losses
        """
        raise NotImplementedError

    def on_train_epoch_start(self):
        self.train_losses = {loss: AverageMeter(loss, ':.2f')
                             for loss in self.get_loss_names()}

    def training_step(self, batch, batch_idx):
        tag = batch['tag']
        images = batch['image']
        targets = batch['target']
        idxs = batch['idx']

        x1 = images[tag == 1]
        x2 = images[tag == 2]
        target = targets[tag == 1]
        idx_ulb = idxs[tag == 2]
        labeled_len = len(target)
        total_len = len(targets)
        x = torch.cat([x1, x2], dim=0)

        output, (f_g, f_p) = self.encoder(x)

        # Calculate the loss using the loss function
        loss_map = self.loss(labeled_len=labeled_len, total_len=total_len, output=output,
                             feat=f_g, feat_p=f_p, target=target, idx_ulb=idx_ulb)
        # Get the total loss from the loss map
        loss = loss_map['total_loss']
        for key, value in loss_map.items():
            update_meter(
                self.train_losses[key], value, self.args.train.batch_size, self.args.distributed)
        for ls in self.train_losses.values():
            self.log(ls.name, ls.avg, on_step=True, prog_bar=True)

        return loss

    def on_train_epoch_end(self):
        results = {key: value.avg for key, value in self.train_losses.items()}
        wandb.log(results, step=self.current_epoch)

    def validation_step(self, batch, batch_idx):
        tags = batch['tag']
        images = batch['image']
        targets = batch['target']

        image = images[tags != 0]
        target = targets[tags != 0]
        tag = tags[tags != 0]

        output, _ = self.encoder(image)
        prob = F.softmax(output, dim=1)
        conf, pred = prob.max(1)

        self.val_step_outputs['tags'].extend(tag)
        self.val_step_outputs['preds'].extend(pred)
        self.val_step_outputs['label'].extend(target)
        self.val_step_outputs['conf'].extend(conf)
        return pred

owdfa/algorithm/ssl.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
from .dataset_factory import create_dataloader

owdfa/datasets/__init__.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
from loguru import logger
import torch.utils.data as data

from owdfa.datasets.dfa import DFA
from owdfa.datasets.utils import create_data_transforms, MultilabelBalancedRandomSampler


def create_dataloader(args, split, loader='torch'):
    if args.dataset.loader == 'torch':
        return create_torch_dataloader(args, split)
    else:
        logger.error(f'Unknown loader: {args.dataset.loader}')


def create_torch_dataloader(args, split):
    num_workers = args.num_workers if 'num_workers' in args else 8
    balance_sample = args.balance_sample if 'balance_sample' in args else False
    batch_size = getattr(args, split).batch_size

    transform = create_data_transforms(args.transform, split)
    kwargs = getattr(args.dataset, args.dataset.name)
    dataset = eval(args.dataset.name)(
        split=split, transform=transform, **kwargs)

    sampler = None

    if args.distributed:
        sampler = data.DistributedSampler(dataset)

    if balance_sample and split == 'train':
        sampler = MultilabelBalancedRandomSampler(dataset)

    if args.debug:
        dataset = data.Subset(dataset, range(0, 100))
        sampler = None
        logger.warning(
            'Distributed dataset sampler is disabled in debug mode!')

    shuffle = True if sampler is None and split == 'train' else False

    dataloader = data.DataLoader(dataset,
                                 batch_size=batch_size,
                                 shuffle=shuffle,
                                 sampler=sampler,
                                 num_workers=num_workers,
                                 pin_memory=True,
                                 )

    return dataloader

owdfa/datasets/dataset_factory.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
import os
import cv2
import pandas as pd

from torch.utils.data import Dataset
from loguru import logger

from owdfa.datasets.utils import dlib_crop_face

try:
    import dlib
except:
    print('Please install dlib when using face detection!')


class DFA(Dataset):
    def __init__(self,
                 data_root,
                 split='train',
                 meta_paths=None,
                 train_ratio_per_class=0.8,
                 seed=1234,
                 transform=None,
                 transform_twice=False,
                 crop_face=True,
                 predictor_path=None,
                 mod='all',
                 pseudo=False,
                 **kwargs
                 ):
        self.data_root = data_root
        self.split = split
        self.meta_paths = meta_paths
        self.train_ratio_per_class = train_ratio_per_class
        self.transform = transform
        self.transform_twice = transform_twice
        self.seed = seed
        self.crop_face = crop_face
        self.predictor_path = predictor_path
        self.mod = mod
        self.pseudo = pseudo

        self.setup()
        self.default_data = None

        if self.crop_face:
            logger.info("Using dlib for face detection!")
            self.detector = dlib.get_frontal_face_detector()
            self.predictor = dlib.shape_predictor(self.predictor_path)

    def setup(self):
        if type(self.meta_paths) is str:
            self.meta_paths = [self.meta_paths]

        train_dfs = []
        test_dfs = []
        for meta_path in self.meta_paths:
            assert os.path.exists(meta_path), f'{meta_path} not exists!'
            sub_df = pd.read_csv(meta_path)
            train_df, test_df = self._split_data(
                sub_df, self.train_ratio_per_class, self.seed)
            train_dfs.append(train_df)
            test_dfs.append(test_df)
        train_dfs = pd.concat(train_dfs, ignore_index=True)
        test_dfs = pd.concat(test_dfs, ignore_index=True)

        if self.pseudo:
            idx = train_dfs[train_dfs.tag == 2].index
            train_dfs.loc[idx, 'label'] = train_dfs.loc[idx, 'pseudo']

        self.items = train_dfs if self.split == 'train' else test_dfs

        logger.info(f'Total {self.split} numbers: {len(self.items)}')

    def _split_data(self, meta_df, train_ratio_per_class, seed):
        if self.mod == 'all':
            meta_df = meta_df[meta_df.tag != 0]
        elif self.mod == 'labeled':
            meta_df = meta_df[meta_df.tag == 1]
        elif self.mod == 'unlabeled':
            meta_df = meta_df[meta_df.tag == 2]
        else:
            assert False

        train_dfs = []
        for source in sorted(meta_df.image_source.unique()):
            for label in sorted(meta_df.label.unique()):
                train_df = meta_df[meta_df.image_source == source][meta_df.label == label].sample(
                    frac=train_ratio_per_class, random_state=seed)
                train_dfs.append(train_df)
        train_df = pd.concat(train_dfs)
        test_df = meta_df[~meta_df.index.isin(train_df.index)]

        train_df.reset_index(drop=True, inplace=True)
        test_df.reset_index(drop=True, inplace=True)
        return train_df, test_df

    def load_img(self, img_path):
        img = cv2.imread(img_path)

        if self.crop_face:
            # margin = random.uniform(1.0, 1.5) if self.split == 'train' else 1.2
            img = dlib_crop_face(img, self.detector,
                                 self.predictor, align=False, margin=1.2)

        img = img[:, :, ::-1]
        return img

    def get_method(self):
        return {label: method for method, label in
                self.items.label.groupby(self.items.method).value_counts().keys()}

    def get_labels(self):
        return self.items.label

    def get_tags(self):
        return self.items.tag

    def __getitem__(self, index):
        img_path, label, face_type, tag = self.items.loc[index, [
            'img_path', 'label', 'face_type', 'tag']]

        try:
            img = self.load_img(img_path)

            if self.transform_twice:
                img1 = self.transform(image=img)['image']
                img2 = self.transform(image=img)['image']
                img = [img1, img2]
            else:
                if self.transform is not None:
                    img = self.transform(image=img)['image']

            item = {'image': img, 'target': label, 'idx': index,
                    'tag': tag, 'face_type': face_type, 'img_path': img_path}

        except Exception as e:
            logger.info(f'{e}: {img_path}')
            return self.default_data

        if self.default_data is None:
            self.default_data = item

        return item

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

    def __exit__(self):
        if self.lmdb_dir is not None:
            self.lmdb.close()

owdfa/datasets/dfa.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
from .transforms import create_data_transforms
from .face_cropper import *
from .sampler import *

owdfa/datasets/utils/__init__.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
try:
    import dlib
except:
    print('Please install dlib when using face detection!')


def add_face_margin(x, y, w, h, margin):
    assert margin >= 1.0, "margin must be greater than 1.0"
    margin -= 1
    x0 = x - int(w * margin / 2)
    y0 = y - int(h * margin / 2)
    x1 = x + w + int(w * margin / 2)
    y1 = y + h + int(h * margin / 2)
    return [x0, y0, x1, y1]


def get_default_bbox(img, face_boxes, margin):
    x0, y0, x1, y1 = face_boxes
    x0, y0, x1, y1 = add_face_margin(x0, y0, x1-x0, y1-y0, margin)
    max_h, max_w = img.shape[:2]
    x0 = int(max(0, x0))
    y0 = int(max(0, y0))
    x1 = int(min(max_w, x1))
    y1 = int(min(max_h, y1))
    return [x0, y0, x1, y1]


def dlib_crop_face(img, detector, predictor, align=False, margin=1.2):
    dets = detector(img, 0)
    if len(dets) > 0:
        det = dets[0]

        if align:
            shape = predictor(img, det)
            face = dlib.get_face_chip(img, shape, size=320)
        else:
            bbox = [det.left(), det.top(), det.right(), det.bottom()]
            x0, y0, x1, y1 = get_default_bbox(img, bbox, margin=margin)
            assert x0 < x1 and y0 < y1
            face = img[y0:y1, x0:x1]

        img = face
    return img

owdfa/datasets/utils/face_cropper.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
import random
import numpy as np

import torch
import torch.utils.data
import torchvision
from torch.utils.data.distributed import DistributedSampler


class MultilabelBalancedRandomSampler(DistributedSampler):
    """
    MultilabelBalancedRandomSampler: Given a multilabel dataset of length n_samples and
    number of classes n_classes, samples from the data with equal probability per class
    effectively oversampling minority classes and undersampling majority classes at the
    same time. Note that using this sampler does not guarantee that the distribution of
    classes in the output samples will be uniform, since the dataset is multilabel and
    sampling is based on a single class. This does however guarantee that all classes
    will have at least batch_size / n_classes samples as batch_size approaches infinity
    """

    def __init__(self, dataset, labels=None, indices=None, class_choice="balance", callback_get_label=None):
        """
        Parameters:
        -----------
            labels: a multi-hot encoding numpy array of shape (n_samples, n_classes)
            indices: an arbitrary-length 1-dimensional numpy array representing a list
            of indices to sample only from
            class_choice: a string indicating how class will be selected for every
            sample:
                "least_sampled": class with the least number of sampled labels so far
                "random": class is chosen uniformly at random
                "cycle": the sampler cycles through the classes sequentially
        """
        self.indices = range(len(dataset)) if indices is None else indices
        self.callback_get_label = callback_get_label

        self.tag = dataset.get_tags()
        self.labels = self._get_labels(dataset) if labels is None else labels
        self.num_classes = np.max(self.labels) + 1
        self.num_label_classes = np.max(self.labels[self.tag == 1]) + 1
        self.labels = np.eye(self.num_classes)[self.labels]
        self.tag = np.eye(3)[self.tag]

        # List of lists of example indices per class
        self.class_indices = []
        # Labeled data
        for class_ in range(0, self.num_label_classes):
            class_mask = self.labels[:, class_] == 1
            seen_mask = self.tag[:, 1] == 1
            lst = np.where(class_mask & seen_mask)[0]
            lst = lst[np.isin(lst, self.indices)]
            self.class_indices.append(np.random.permutation(lst))
        # Unlabeled data
        unseen_mask = self.tag[:, 2] == 1
        lst = np.where(unseen_mask)[0]
        lst = lst[np.isin(lst, self.indices)]
        self.class_indices.append(np.random.permutation(lst))

        self.counts = [0] * (self.num_label_classes + 1)
        self.stat = []

        assert class_choice in ["random", "cycle", "balance"]
        self.class_choice = class_choice
        self.balance = True

    def __iter__(self):
        self.count = 0
        return self

    def __next__(self):
        if self.count >= len(self.indices):
            raise StopIteration
        self.count += 1
        return self.sample()

    def sample(self):
        class_ = self.get_class()
        class_indices = self.class_indices[class_]
        chosen_index = class_indices[self.counts[class_] % len(class_indices)]
        self.counts[class_] += 1
        return chosen_index

    def get_class(self):
        if self.class_choice == "random":
            class_ = random.randint(0, self.labels.shape[1] - 1)
        elif self.class_choice == "cycle":
            class_ = self.current_class
            self.current_class = (self.current_class + 1) % self.labels.shape[1]
        elif self.class_choice == "balance":
            if self.balance:
                class_ = self.num_label_classes
            else:
                min_count = self.counts[0]
                min_classes = [0]
                for class_ in range(1, self.num_label_classes):
                    if self.counts[class_] < min_count:
                        min_count = self.counts[class_]
                        min_classes = [class_]
                    if self.counts[class_] == min_count:
                        min_classes.append(class_)
                class_ = random.choice(min_classes)
            self.balance = not self.balance
        return class_

    def _get_labels(self, dataset):
        if self.callback_get_label:
            return self.callback_get_label(dataset)
        elif isinstance(dataset, torch.utils.data.TensorDataset):
            return dataset.tensors[1]
        elif isinstance(dataset, torchvision.datasets.MNIST):
            return dataset.train_labels.tolist()
        elif isinstance(dataset, torchvision.datasets.ImageFolder):
            return [x[1] for x in dataset.imgs]
        elif isinstance(dataset, torchvision.datasets.DatasetFolder):
            return dataset.samples[:][1]
        elif isinstance(dataset, torch.utils.data.Subset):
            return dataset.dataset.imgs[:][1]
        elif isinstance(dataset, torch.utils.data.Dataset):
            return dataset.get_labels()
        else:
            raise NotImplementedError

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

    def set_epoch(self, epoch):
        pass

owdfa/datasets/utils/sampler.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
import albumentations as alb
from albumentations.pytorch.transforms import ToTensorV2


def create_data_transforms(args, split='train'):
    image_size = getattr(args, 'image_size', 224)
    mean = getattr(args, 'mean', [0.485, 0.456, 0.406])
    std = getattr(args, 'std', [0.229, 0.224, 0.225])
    num_segments = args.num_segments if 'num_segments' in args else 1

    additional_targets = {}
    for i in range(1, num_segments):
        additional_targets[f'image{i}'] = 'image'

    if split == 'train':
        transfrom = alb.Compose([
            alb.HorizontalFlip(),
            alb.RandomResizedCrop(int((256 / 224) * image_size),
                                  int((256 / 224) * image_size), scale=(0.5, 1.0), p=0.2),
            alb.RandomBrightnessContrast(p=0.2),
            alb.Resize(image_size, image_size),
            alb.Normalize(mean=mean, std=std),
            ToTensorV2(),
        ], additional_targets=additional_targets)
    else:
        transfrom = alb.Compose([
            alb.Resize(image_size, image_size),
            alb.Normalize(mean=mean, std=std),
            ToTensorV2(),
        ], additional_targets=additional_targets)

    return transfrom

owdfa/datasets/utils/transforms.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
from .base_net import BaseClassifier, CPLClassifier

owdfa/encoder/__init__.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
import warnings
import numpy as np
import timm
from loguru import logger

import torch
import torch.nn as nn


__all__ = ['BinaryClassifier']

MODEL_DICTS = {}
MODEL_DICTS.update(timm.models.__dict__)

warnings.filterwarnings("ignore")


def weights_init_kaiming(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        nn.init.kaiming_normal_(m.weight, a=0, mode='fan_out')
        nn.init.constant_(m.bias, 0.0)
    elif classname.find('Conv') != -1:
        nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in')
        if m.bias is not None:
            nn.init.constant_(m.bias, 0.0)
    elif classname.find('BatchNorm') != -1:
        if m.affine:
            nn.init.constant_(m.weight, 1.0)
            nn.init.constant_(m.bias, 0.0)


def weights_init_classifier(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        nn.init.normal_(m.weight, std=0.001)
        if m.bias:
            nn.init.constant_(m.bias, 0.0)


class AdaptiveAvgPool2dCustom(nn.Module):
    def __init__(self, output_size):
        super(AdaptiveAvgPool2dCustom, self).__init__()
        self.output_size = np.array(output_size)

    def forward(self, x: torch.Tensor):
        stride_size = np.floor(
            np.array(x.shape[-2:]) / self.output_size).astype(np.int32)
        kernel_size = np.array(x.shape[-2:]) - \
            (self.output_size - 1) * stride_size
        avg = nn.AvgPool2d(kernel_size=list(kernel_size),
                           stride=list(stride_size))
        x = avg(x)
        return x


class BaseClassifier(nn.Module):
    def __init__(self,
                 encoder,
                 num_classes=20,
                 drop_rate=0.2,
                 is_feat=False,
                 neck='no',
                 pretrained=False,
                 **kwargs) -> None:
        super().__init__()
        self.num_classes = num_classes
        self.drop_rate = drop_rate
        self.is_feat = is_feat
        self.neck = neck

        self.encoder = MODEL_DICTS[encoder](pretrained=pretrained, **kwargs)
        for k, v in kwargs.items():
            setattr(self, k, v)

        if hasattr(self.encoder, 'get_classifier'):
            self.num_features = self.encoder.get_classifier().in_features
        else:
            self.num_features = self.encoder.last_channel

        # self.pool = nn.AdaptiveAvgPool2d((1, 1))

        self.pool = AdaptiveAvgPool2dCustom((1, 1))

        self.dropout = nn.Dropout(drop_rate)

        if self.neck == 'bnneck':
            logger.info('Using BNNeck')
            self.bottleneck = nn.BatchNorm1d(self.num_features)
            self.bottleneck.bias.requires_grad_(False)  # no shift
            self.fc2 = nn.Linear(
                self.num_features, self.num_classes, bias=False)
            self.bottleneck.apply(weights_init_kaiming)
            self.fc2.apply(weights_init_classifier)
        else:
            self.fc2 = nn.Linear(self.num_features, self.num_classes)

    def forward_featuremaps(self, x):
        featuremap = self.encoder.forward_features(x)
        return featuremap

    def forward_features(self, x):
        featuremap = self.encoder.forward_features(x)
        feature = self.pool(featuremap).flatten(1)

        if self.neck == 'bnneck':
            feature = self.bottleneck(feature)

        return feature

    def forward(self, x, label=None):
        feature = self.forward_features(x)

        x = self.dropout(feature)
        method = self.fc2(x)

        y = method

        if self.is_feat:
            return y, feature

        return y


class CPLClassifier(nn.Module):
    def __init__(self,
                 encoder,
                 num_classes=20,
                 num_patch=3,
                 drop_rate=0.2,
                 is_feat=False,
                 neck='bnneck',
                 pretrained=False,
                 **kwargs) -> None:
        super().__init__()
        self.num_classes = num_classes
        self.num_patch = num_patch
        self.drop_rate = drop_rate
        self.is_feat = is_feat
        self.neck = neck

        self.encoder = MODEL_DICTS[encoder](pretrained=pretrained, **kwargs)
        for k, v in kwargs.items():
            setattr(self, k, v)

        if hasattr(self.encoder, 'get_classifier'):
            self.num_features = self.encoder.get_classifier().in_features
        else:
            self.num_features = self.encoder.last_channel

        # self.pool = nn.AdaptiveAvgPool2d((1, 1))
        # self.part = nn.AdaptiveAvgPool2d((self.num_patch, self.num_patch))

        self.pool = AdaptiveAvgPool2dCustom((1, 1))
        self.part = AdaptiveAvgPool2dCustom((self.num_patch, self.num_patch))

        self.dropout = nn.Dropout(drop_rate)

        self.bottleneck = nn.BatchNorm1d(self.num_features)
        self.bottleneck.bias.requires_grad_(False)  # no shift
        self.fc2 = nn.Linear(self.num_features, self.num_classes, bias=False)
        self.bottleneck.apply(weights_init_kaiming)
        self.fc2.apply(weights_init_classifier)

        for i in range(self.num_patch ** 2):
            name = 'bnneck' + str(i)
            setattr(self, name, nn.BatchNorm1d(self.num_features))
            getattr(self, name).bias.requires_grad_(False)
            getattr(self, name).apply(weights_init_kaiming)

    def forward_featuremaps(self, x):
        featuremap = self.encoder.forward_features(x)
        return featuremap

    def forward_features(self, x):
        featuremap = self.encoder.forward_features(x)
        f_g = self.pool(featuremap).flatten(1)
        f_p = self.part(featuremap)
        f_p = f_p.view(f_p.size(0), f_p.size(1), -1)

        fs_p = []
        for i in range(self.num_patch ** 2):
            f_p_i = f_p[:, :, i]
            f_p_i = getattr(self, 'bnneck' + str(i))(f_p_i)
            fs_p.append(f_p_i)
        fs_p = torch.stack(fs_p, dim=-1)

        f_g = self.bottleneck(f_g)

        return (f_g, fs_p)

    def forward(self, x, label=None):
        feature = self.forward_features(x)
        (f_g, _) = feature

        x = self.dropout(f_g)
        method = self.fc2(x)

        y = method

        if self.is_feat:
            return y, feature

        return y

owdfa/encoder/base_net.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
from .label_smooth import LabelSmoothSoftmaxCEV1
from .losses import entropy, symmetric_mse_loss, MarginLoss

owdfa/losses/__init__.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
#!/usr/bin/python
# -*- encoding: utf-8 -*-
import torch
import torch.nn as nn


class LabelSmoothSoftmaxCEV1(nn.Module):
    '''
    This is the autograd version, you can also try the LabelSmoothSoftmaxCEV2 that uses derived gradients
    '''

    def __init__(self, lb_smooth=0.1, reduction='mean', ignore_index=-100):
        super(LabelSmoothSoftmaxCEV1, self).__init__()
        self.lb_smooth = lb_smooth
        self.reduction = reduction
        self.lb_ignore = ignore_index
        self.log_softmax = nn.LogSoftmax(dim=1)

    def forward(self, logits, label):
        '''
        Same usage method as nn.CrossEntropyLoss:
            >>> criteria = LabelSmoothSoftmaxCEV1()
            >>> logits = torch.randn(8, 19, 384, 384) # nchw, float/half
            >>> lbs = torch.randint(0, 19, (8, 384, 384)) # nhw, int64_t
            >>> loss = criteria(logits, lbs)
        '''
        # overcome ignored label
        logits = logits.float()  # use fp32 to avoid nan
        with torch.no_grad():
            num_classes = logits.size(1)
            label = label.clone().detach()
            ignore = label.eq(self.lb_ignore)
            n_valid = ignore.eq(0).sum()
            label[ignore] = 0
            lb_pos, lb_neg = 1. - self.lb_smooth, self.lb_smooth / num_classes
            lb_one_hot = torch.empty_like(logits).fill_(
                lb_neg).scatter_(1, label.unsqueeze(1), lb_pos).detach()

        logs = self.log_softmax(logits)
        loss = -torch.sum(logs * lb_one_hot, dim=1)
        loss[ignore] = 0
        if self.reduction == 'mean':
            loss = loss.sum() / n_valid
        if self.reduction == 'sum':
            loss = loss.sum()

        return loss

owdfa/losses/label_smooth.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
#!/usr/bin/python
# -*- encoding: utf-8 -*-
import torch
import torch.nn as nn
import torch.nn.functional as F


class MarginLoss(nn.Module):
    def __init__(self, m=0.2, weight=None, s=10):
        super(MarginLoss, self).__init__()
        self.m = m
        self.s = s
        self.weight = weight

    def forward(self, x, target):
        index = torch.zeros_like(x, dtype=torch.uint8)
        index.scatter_(1, target.data.view(-1, 1), 1)
        x_m = x - self.m * self.s
        output = torch.where(index, x_m, x)
        return F.cross_entropy(output, target, weight=self.weight)


def entropy(x):
    """
    Helper function to compute the entropy over the batch
    input: batch w/ shape [b, num_classes]
    output: entropy value [is ideally -log(num_classes)]
    """
    eps = 1e-8
    x_ = torch.clamp(x, min=eps)
    b = x_ * torch.log(x_)

    if len(b.size()) == 2:  # Sample-wise entropy
        return - b.sum(dim=1).mean()
    elif len(b.size()) == 1:  # Distribution-wise entropy
        return - b.sum()
    else:
        raise ValueError('Input tensor is %d-Dimensional' % (len(b.size())))


def symmetric_mse_loss(input1, input2):
    """
    Like F.mse_loss but sends gradients to both directions

    Note:
    - Returns the sum over all examples. Divide by the batch size afterwards
      if you want the mean.
    - Sends gradients to both input1 and input2.
    """
    assert input1.size() == input2.size()
    return torch.sum((input1 - input2)**2)

owdfa/losses/losses.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
from torch.optim import *
from torch_optimizer import *

owdfa/optimizers/__init__.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
from .parameters import get_parameters
from .meters import AverageMeter, ProgressMeter
from .vis import visualize
from .eval import val_stat
from .misc import *

owdfa/utils/__init__.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
#!/usr/bin/env python
# coding: utf-8
import numpy as np
from loguru import logger
from scipy.optimize import linear_sum_assignment
from sklearn.metrics import normalized_mutual_info_score, adjusted_rand_score


def direct_acc(y_pred, y_true):
    """
    Calculate accuracy
    # Arguments
        y_pred: predicted labels, numpy.array with shape `(n_samples,)`
        y_true: true labels, numpy.array with shape `(n_samples,)`
    # Return
        accuracy, in [0,1]
    """
    num_correct = np.sum(y_pred == y_true)
    res = num_correct / len(y_true)
    return res


def cluster_map(y_pred, y_true):
    """
    Calculate clustering mapping relation. Require scipy installed
    # Arguments
        y_pred: predicted labels, numpy.array with shape `(n_samples,)`
        y_true: true labels, numpy.array with shape `(n_samples,)`
    # Return
        row_ind, col_ind
    """
    y_true = y_true.astype(np.int64)
    assert y_pred.size == y_true.size
    d = max(y_pred.max(), y_true.max()) + 1
    w = np.zeros((d, d), dtype=np.int64)
    for i in range(y_pred.size):
        w[y_pred[i], y_true[i]] += 1
    _, col_ind = linear_sum_assignment(w.max() - w)
    row_ind = {col_ind[x]: x for x in col_ind}
    return row_ind, col_ind


def cluster_acc(y_pred, y_true):
    """
    Calculate clustering accuracy. Require scipy installed
    # Arguments
        y_pred: predicted labels, numpy.array with shape `(n_samples,)`
        y_true: true labels, numpy.array with shape `(n_samples,)`
    # Return
        accuracy, in [0,1]
    """
    y_true = y_true.astype(np.int64)
    assert y_pred.size == y_true.size
    d = max(y_pred.max(), y_true.max()) + 1
    w = np.zeros((d, d), dtype=np.int64)
    for i in range(y_pred.size):
        w[y_pred[i], y_true[i]] += 1
    row_ind, col_ind = linear_sum_assignment(w.max() - w)
    return w[row_ind, col_ind].sum() / y_pred.size


def nmi(y_pred, y_true):
    """
    Calculate normalized mutual information. Require scikit-learn installed
    # Arguments
        y_pred: predicted labels, numpy.array with shape `(n_samples,)`
        y_true: true labels, numpy.array with shape `(n_samples,)`
    # Return
        NMI, in [0,1]
    """
    return normalized_mutual_info_score(y_true, y_pred)


def ari(y_pred, y_true):
    """
    Calculate adjusted Rand index. Require scikit-learn installed
    # Arguments
        y_pred: predicted labels, numpy.array with shape `(n_samples,)`
        y_true: true labels, numpy.array with shape `(n_samples,)`
    # Return
        ARI, in [-1,1]
    """
    return adjusted_rand_score(y_true, y_pred)


def val_stat(y_tags, y_pred, y_label, y_conf):
    logger.info(f'##### Eval Results #####')
    num_known = np.max(y_label[y_tags == 1], axis=0) + 1
    known_class = y_label < num_known
    novel_class = y_label >= num_known

    mean_uncert = 1 - np.mean(y_conf)
    over_acc = cluster_acc(y_pred, y_label)
    over_nmi = nmi(y_pred, y_label)
    over_ari = ari(y_pred, y_label)
    known_acc = direct_acc(y_pred[known_class], y_label[known_class])
    known_nmi = nmi(y_pred[known_class], y_label[known_class])
    known_ari = ari(y_pred[known_class], y_label[known_class])
    if len(y_pred[novel_class]) > 0:
        novel_acc = cluster_acc(y_pred[novel_class], y_label[novel_class])
        novel_nmi = nmi(y_pred[novel_class], y_label[novel_class])
        novel_ari = ari(y_pred[novel_class], y_label[novel_class])
    else:
        novel_acc = novel_nmi = novel_ari = 0

    results = {
        'all_acc': over_acc,
        'all_nmi': over_nmi,
        'all_ari': over_ari,
        'known_acc': known_acc,
        'novel_acc': novel_acc,
        'known_nmi': known_nmi,
        'novel_nmi': novel_nmi,
        'known_ari': known_ari,
        'novel_ari': novel_ari,
        'val_uncert': mean_uncert,
    }

    logger.info(f'Uncertainty: {mean_uncert:.4f}')
    logger.info(
        f'[All] ACC: {over_acc:.4f}, NMI: {over_nmi:.4f}, ARI: {over_ari:.4f}')
    logger.info(
        f'[Known] ACC: {known_acc:.4f}, NMI: {known_nmi:.4f}, ARI: {known_ari:.4f}')
    logger.info(
        f'[Novel] ACC: {novel_acc:.4f}, NMI: {novel_nmi:.4f}, ARI: {novel_ari:.4f}')
    logger.info(f'{known_acc * 100:.2f}, '
                + f'{novel_acc * 100:.2f}, {novel_nmi * 100:.2f}, {novel_ari * 100:.2f}, '
                + f'{over_acc * 100:.2f}, {over_nmi * 100:.2f}, {over_ari * 100:.2f}')

    return results

owdfa/utils/eval.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
#!/usr/bin/env python
# coding: utf-8
class AverageMeter(object):
    """Computes and stores the average and current value.

       Code imported from https://github.com/pytorch/examples/blob/master/imagenet/main.py#L247-L262
    """

    def __init__(self, name='metric', fmt=':f'):
        self.name = name
        self.fmt = fmt
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

    def __str__(self):
        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
        return fmtstr.format(**self.__dict__)


class ProgressMeter(object):
    def __init__(self, num_batches, meters, prefix=""):
        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)
        self.meters = meters
        self.prefix = prefix

    def display(self, batch):
        entries = [self.prefix + self.batch_fmtstr.format(batch)]
        entries += [str(meter) for meter in self.meters]
        # print('\t'.join(entries))
        return '\t'.join(entries)

    def _get_batch_fmtstr(self, num_batches):
        num_digits = len(str(num_batches // 1))
        fmt = '{:' + str(num_digits) + 'd}'
        return '[' + fmt + '/' + fmt.format(num_batches) + ']'

owdfa/utils/meters.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
#!/usr/bin/env python
# coding: utf-8
import os
import random
import shutil
import warnings
import wandb
import numpy as np
import torch
import torch.distributed as dist


def set_seed(args):
    seed = args.seed
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


def setup(args):
    os.environ['TORCH_HOME'] = args.torch_home
    warnings.filterwarnings("ignore")
    set_seed(args)


def init_wandb_workspace(args):
    if args.wandb.name is None:
        args.wandb.name = args.config.split('/')[-1].replace('.yaml', '')

    if args.local_rank == 0:
        wandb.init(**args.wandb)
        allow_val_change = False if args.wandb.resume is None else True
        wandb.config.update(dict(args), allow_val_change)
        wandb.save(args.config)

        if args.debug or wandb.run.dir == '/tmp':
            args.exam_dir = 'wandb/debug'
        else:
            args.exam_dir = os.path.dirname(wandb.run.dir)

        if os.path.exists(args.exam_dir):
            shutil.rmtree(args.exam_dir)
        os.makedirs(args.exam_dir, exist_ok=True)
        os.makedirs(os.path.join(args.exam_dir, 'ckpts'), exist_ok=True)

    return args


def save_test_results(img_paths, y_preds, y_trues, filename='results.log'):
    assert len(y_trues) == len(y_preds) == len(img_paths)

    with open(filename, 'w') as f:
        for i in range(len(img_paths)):
            print(img_paths[i], end=' ', file=f)
            print(y_preds[i], file=f)
            print(y_trues[i], end=' ', file=f)


def gather_tensor(inp, world_size=None, dist_=True, to_numpy=False):
    """Gather tensor in the distributed setting.

    Args:
        inp (torch.tensor): 
            Input torch tensor to gather.
        world_size (int, optional): 
            Dist world size. Defaults to None. If None, world_size = dist.get_world_size().
        dist_ (bool, optional):
            Whether to use all_gather method to gather all the tensors. Defaults to True.
        to_numpy (bool, optional): 
            Whether to return numpy array. Defaults to False.

    Returns:
        (torch.tensor || numpy.ndarray): Returned tensor or numpy array.
    """
    inp = torch.stack(inp)
    if dist_:
        if world_size is None:
            world_size = dist.get_world_size()
        gather_inp = [torch.ones_like(inp) for _ in range(world_size)]
        dist.all_gather(gather_inp, inp)
        gather_inp = torch.cat(gather_inp)
    else:
        gather_inp = inp

    if to_numpy:
        gather_inp = gather_inp.cpu().numpy()

    return gather_inp


@torch.no_grad()
def concat_all_gather(tensor: torch.Tensor) -> torch.Tensor:
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    tensors_gather = [torch.ones_like(tensor)
                      for _ in range(torch.distributed.get_world_size())]
    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)

    output = torch.cat(tensors_gather, dim=0)
    return output


@torch.no_grad()
def reduce_tensor(tensor: torch.Tensor) -> torch.Tensor:
    rt = tensor.clone()
    torch.distributed.all_reduce(rt, op=torch.distributed.ReduceOp.SUM)
    rt /= torch.distributed.get_world_size()
    return rt


def update_meter(meter, value, size, is_dist=False):
    if is_dist:
        meter.update(reduce_tensor(value.data).item(), size)
    else:
        meter.update(value.item(), size)
    return meter

owdfa/utils/misc.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
#!/usr/bin/env python
# coding: utf-8
import os
import argparse
import omegaconf
from loguru import logger


def flatten_dict(dictionary, exclude=None, delimiter='.'):
    if exclude is None:
        exclude = []
    flat_dict = dict()
    for key, value in dictionary.items():
        if isinstance(value, dict) and key not in exclude:
            flatten_value_dict = flatten_dict(value, exclude, delimiter)
            for k, v in flatten_value_dict.items():
                flat_dict[f"{key}{delimiter}{k}"] = v
        else:
            flat_dict[key] = value
    return flat_dict


def nested_set(dic, keys, value):
    for key in keys[:-1]:
        dic = dic.setdefault(key, {})
    dic[keys[-1]] = value


def get_parameters(to_dict=False, config=None, **new_kwargs):
    # priority: cmd args > new_kwargs > dict in config
    # `--config` in cmd args can overwrite `config`

    # (1) the base parser, introduce `config`
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument('-c', '--config', type=str,
                        default=config or 'cfg.yaml')
    parser.add_argument('--local_rank', type=int, default=0)
    parser.add_argument('--world_size', type=int, default=1)
    parser.add_argument('--sync-bn', action='store_true', default=False)
    parser.add_argument('--debug', action='store_true', default=False)
    # parse the above cmd options
    args_tmp = parser.parse_known_args()[0]
    args_tmp.distributed = False
    if 'WORLD_SIZE' in os.environ:
        args_tmp.distributed = int(os.environ['WORLD_SIZE']) > 1
    if args_tmp.debug:
        args_tmp.train.epochs = 2
    args_tmp_dict = vars(args_tmp)

    try:
        oc_cfg = omegaconf.OmegaConf.load(args_tmp.config)
    except Exception as e:
        logger.warning(f'config path [{args_tmp.config}] is not valid.')
        raise e
    oc_cfg.merge_with(args_tmp_dict)

    # (2) append items from new_kwargs, new_kwargs has higher priority
    if new_kwargs:
        for k in new_kwargs:
            if k in oc_cfg:
                logger.warning(
                    f'{k} from `new_kwargs` found in original conf, will keep the one in `new_kwargs`')
        oc_cfg.merge_with(omegaconf.OmegaConf.create(new_kwargs))

    oc_cfg_dict = omegaconf.OmegaConf.to_container(oc_cfg, resolve=True)
    # turn chained dict to flatten dict, turn {a: {b: {c: 1}}} to a.b.c = 1
    oc_cfg_dict_flatten = flatten_dict(oc_cfg_dict)

    # (3) add all the k-v items from `oc_cfg_dict_flatten` to argparse
    # dest change dotted names to be connected with '___', as dotted dest name is not supported by argparse
    for k, v in oc_cfg_dict_flatten.items():
        if k in args_tmp_dict:
            continue
        # not using the default store_true syntax
        # use lambda func so that we can use `python a.py --switch True`
        if isinstance(v, bool):
            parser.add_argument('--{}'.format(k), dest=k.replace('.', '___'),
                                type=lambda x: (str(x).lower() == 'true'), default=v)
        elif isinstance(v, list) or isinstance(v, tuple):
            parser.add_argument(
                '--{}'.format(k), dest=k.replace('.', '___'), type=type(v[0]), default=v, nargs='+')
        else:
            # None type default changed to str
            parser.add_argument('--{}'.format(k), dest=k.replace('.', '___'),
                                type=str if v is None else type(v), default=v)
    parser.add_argument('-h', '--help', action='help',
                        help=('show this help message and exit'))
    args = parser.parse_args()

    var_args = vars(args)
    for k, v in var_args.items():
        ori_k = k.replace('___', '.')
        sub_ks = ori_k.split('.')
        # turn back to nested dict
        # turn a.b.c = 1 to {a: {b: {c: 1}}}
        nested_set(oc_cfg_dict, sub_ks, v)

    oc_cfg.merge_with(oc_cfg_dict)

    # print(OmegaConf.to_yaml(oc_cfg))
    if to_dict:
        return omegaconf.OmegaConf.to_container(oc_cfg, resolve=True)
    return oc_cfg


if __name__ == '__main__':
    args = get_parameters()
    print(args)

owdfa/utils/parameters.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
#!/usr/bin/env python
# coding: utf-8
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg
from loguru import logger

from openTSNE import TSNE


MACOSKO_COLORS = {
    "Amacrine cells": "#A5C93D",
    "Astrocytes": "#8B006B",
    "Bipolar cells": "#2000D7",
    "Cones": "#538CBA",
    "Fibroblasts": "#8B006B",
    "Horizontal cells": "#B33B19",
    "Microglia": "#8B006B",
    "Muller glia": "#8B006B",
    "Pericytes": "#8B006B",
    "Retinal ganglion cells": "#C38A1F",
    "Rods": "#538CBA",
    "Vascular endothelium": "#8B006B",
}
ZEISEL_COLORS = {
    "Astroependymal cells": "#d7abd4",
    "Cerebellum neurons": "#2d74bf",
    "Cholinergic, monoaminergic and peptidergic neurons": "#9e3d1b",
    "Di- and mesencephalon neurons": "#3b1b59",
    "Enteric neurons": "#1b5d2f",
    "Hindbrain neurons": "#51bc4c",
    "Immature neural": "#ffcb9a",
    "Immune cells": "#768281",
    "Neural crest-like glia": "#a0daaa",
    "Oligodendrocytes": "#8c7d2b",
    "Peripheral sensory neurons": "#98cc41",
    "Spinal cord neurons": "#c52d94",
    "Sympathetic neurons": "#11337d",
    "Telencephalon interneurons": "#ff9f2b",
    "Telencephalon projecting neurons": "#fea7c1",
    "Vascular cells": "#3d672d",
}
MOUSE_10X_COLORS = {
    0: "#FFFF00",
    1: "#1CE6FF",
    2: "#FF34FF",
    3: "#FF4A46",
    4: "#008941",
    5: "#006FA6",
    6: "#A30059",
    7: "#FFDBE5",
    8: "#7A4900",
    9: "#0000A6",
    10: "#63FFAC",
    11: "#B79762",
    12: "#004D43",
    13: "#8FB0FF",
    14: "#997D87",
    15: "#5A0007",
    16: "#809693",
    17: "#FEFFE6",
    18: "#1B4400",
    19: "#4FC601",
    20: "#3B5DFF",
    21: "#4A3B53",
    22: "#FF2F80",
    23: "#61615A",
    24: "#BA0900",
    25: "#6B7900",
    26: "#00C2A0",
    27: "#FFAA92",
    28: "#FF90C9",
    29: "#B903AA",
    30: "#D16100",
    31: "#DDEFFF",
    32: "#000035",
    33: "#7B4F4B",
    34: "#A1C299",
    35: "#300018",
    36: "#0AA6D8",
    37: "#013349",
    38: "#00846F",
}


def plot(
    x,
    y,
    ax=None,
    title=None,
    draw_legend=True,
    draw_centers=False,
    draw_cluster_labels=False,
    colors=None,
    legend_kwargs=None,
    label_order=None,
    **kwargs
):
    import matplotlib

    if ax is None:
        _, ax = matplotlib.pyplot.subplots(figsize=(12, 12))

    if title is not None:
        ax.set_title(title)

    plot_params = {"alpha": kwargs.get("alpha", 0.8), "s": kwargs.get("s", 1)}

    # Create main plot
    if label_order is not None:
        assert all(np.isin(np.unique(y), label_order))
        classes = [l for l in label_order if l in np.unique(y)]
    else:
        classes = np.unique(y)
    if colors is None:
        default_colors = matplotlib.rcParams["axes.prop_cycle"]
        colors = {k: v["color"] for k, v in zip(classes, default_colors())}

    point_colors = list(map(colors.get, y))

    ax.scatter(x[:, 0], x[:, 1], c=point_colors,
               rasterized=True, **plot_params)

    # Plot mediods
    if draw_centers:
        centers = []
        for yi in classes:
            mask = yi == y
            centers.append(np.median(x[mask, :2], axis=0))
        centers = np.array(centers)

        center_colors = list(map(colors.get, classes))
        ax.scatter(
            centers[:, 0], centers[:, 1], c=center_colors, s=48, alpha=1, edgecolor="k"
        )

        # Draw mediod labels
        if draw_cluster_labels:
            for idx, label in enumerate(classes):
                ax.text(
                    centers[idx, 0],
                    centers[idx, 1] + 2.2,
                    label,
                    fontsize=kwargs.get("fontsize", 6),
                    horizontalalignment="center",
                )

    # Hide ticks and axis
    # ax.set_xticks([]), ax.set_yticks([]), ax.axis("off")

    if draw_legend:
        legend_handles = [
            matplotlib.lines.Line2D(
                [],
                [],
                marker="s",
                color="w",
                markerfacecolor=colors[yi],
                ms=10,
                alpha=1,
                linewidth=0,
                label=yi,
                markeredgecolor="k",
            )
            for yi in classes
        ]
        legend_kwargs_ = dict(
            loc="center left", bbox_to_anchor=(1, 0.5), frameon=False, )
        if legend_kwargs is not None:
            legend_kwargs_.update(legend_kwargs)
        ax.legend(handles=legend_handles, **legend_kwargs_)

    canvas = FigureCanvasAgg(plt.gcf())
    canvas.draw()
    return np.array(canvas.buffer_rgba())


def feature_tsne(x, y):
    logger.info(f'Info of data: {x.shape}, {y.shape}')
    logger.info('TSNE working...')
    x_tsne = TSNE(n_components=2, n_jobs=8).fit(x)
    return plot(x_tsne, y, draw_centers=True, draw_cluster_labels=True, colors=MOUSE_10X_COLORS)


def visualize(tags, labels, feats):
    logger.info(f'Shape of feature and labels: {feats.shape}, {labels.shape}')

    num_known = np.max(labels[tags == 1], axis=0) + 1
    known_mask = labels < num_known
    novel_mask = ~known_mask

    # Known Attack
    known_idxs = np.where(known_mask)[0]
    known_feats = feats[known_idxs]
    known_labels = labels[known_idxs]
    fig1 = feature_tsne(known_feats, known_labels)

    # Novel Attack
    novel_idxs = np.where(novel_mask)[0]
    novel_feats = feats[novel_idxs]
    novel_labels = labels[novel_idxs]
    fig2 = feature_tsne(novel_feats, novel_labels)

    # All Attack
    fig3 = feature_tsne(feats, labels)

    return fig1, fig2, fig3

owdfa/utils/vis.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
import os
import wandb
from loguru import logger
from omegaconf import OmegaConf

import torch
import torch.distributed as dist
import lightning.pytorch as pl
from lightning.pytorch.callbacks import ModelCheckpoint

import owdfa.algorithm as algorithm
from owdfa.datasets import create_dataloader
from owdfa.utils import get_parameters, init_wandb_workspace, setup

import better_exceptions
better_exceptions.hook()

args = get_parameters()
args = init_wandb_workspace(args)
if args.local_rank == 0:
    logger.add(f'{args.exam_dir}/train.log', level="INFO")
    logger.info(OmegaConf.to_yaml(args))


def main():
    # Distributed traning
    if args.distributed:
        args.local_rank = int(os.environ.get('LOCAL_RANK', 0))
        dist.init_process_group(backend='nccl', init_method="env://")
        torch.cuda.set_device(args.local_rank)
        args.world_size = dist.get_world_size()

    # Init setup
    setup(args)

    # Create dataloader
    train_dataloader = create_dataloader(args, split='train')
    test_dataloader = create_dataloader(args, split='test')

    # Resume from checkpoint
    checkpoint_dir = os.path.join(args.exam_dir, 'ckpts') if args.exam_dir else None

    model = algorithm.__dict__[args.method.name](args)

    trainer = pl.Trainer(
        strategy='ddp_find_unused_parameters_true',
        min_epochs=1,
        max_epochs=args.train.epochs,
        default_root_dir=args.exam_dir,
        callbacks=[ModelCheckpoint(
            dirpath=checkpoint_dir,
            verbose=True,
            monitor='all_nmi',
            mode='max',
            save_on_train_epoch_end=True,
        )],
        num_sanity_val_steps=1,
        log_every_n_steps=10,
    )
    trainer.fit(model, train_dataloader, test_dataloader)

    if args.local_rank == 0:
        wandb.finish()


if __name__ == '__main__':
    main()

train.py
TencentYoutuResearch/OpenWorld-DeepFakeAttribution