Official Repository for "SingFake: Singing Voice Deepfake Detection"
import openai, json
import pandas as pd
from tqdm import tqdm

openai.api_key = None # fill in your API key
assert openai.api_key is not None, "You must provide an API key"
systemPrompt = open('systemPrompt.txt', 'r').read().strip()
csv_path = None # fill in the path to the csv file you want to clean
assert csv_path is not None, "You must provide the path to the csv file you want to clean"

csv_file = pd.read_csv(csv_path)

for i in tqdm(range(len(csv_file)), desc="Cleaning"):
    # get url
    url = csv_file.iloc[i]['Url'].strip()
    singer_name = csv_file.iloc[i]['Singer']
    title = csv_file.iloc[i]['Title'].strip()
    try:
        HTMLContent = None # Implement HTMLContent.
        assert HTMLContent is not None, "You must provide an HTML content. This part of the script need to be implemented by yourself."

        response = openai.ChatCompletion.create(
        model="gpt-4",
        messages=[
                {"role": "system", "content": systemPrompt},
                {"role": "user", "content": HTMLContent},
            ]
        )
        # get response text
        response_json = json.loads(response['choices'][0]['message']['content'])

        cleaned_spoof_tag = response_json["type"]
        cleaned_singer_tag = response_json["singer"]
        cleaned_model_tag = response_json["model"]
        cleaned_correctness = str(response_json["correct"])

        
        csv_file.iloc[i]['Bonafide Or Spoof'] = cleaned_spoof_tag
        csv_file.iloc[i]['Singer'] = cleaned_singer_tag
        csv_file.iloc[i]['Model'] = cleaned_model_tag
        if cleaned_correctness.strip().lower() != "true":
            print("WARNING: Correctness is not True at {}".format(csv_file.iloc[i]['Url']))
            print("Singer wrote as {}, but should be {}".format(singer_name, cleaned_singer_tag))
    except Exception as e:
        print(title, e)
        continue
dataset/GPT4-checking/clean_with_gpt.py
yongyizang/SingFake
import os, sys
from tqdm import tqdm
import demucs.separate
from pyannote.audio import Model, Inference

authtoken = None # Fill in your auth token
assert authtoken is not None, "You must provide an auth token to use PyAnnote VAD pipeline."

model = Model.from_pretrained("pyannote/segmentation", use_auth_token=authtoken)

from pyannote.audio.pipelines import VoiceActivityDetection
pipeline = VoiceActivityDetection(segmentation=model)

HYPER_PARAMETERS = {
  # onset/offset activation thresholds
  "onset": 0.5, "offset": 0.5,
  # remove speech regions shorter than that many seconds.
  "min_duration_on": 3.0,
  # fill non-speech regions shorter than that many seconds.
  "min_duration_off": 0.0,
}
pipeline.instantiate(HYPER_PARAMETERS)

def run_vad(file_path):
    vad = pipeline(file_path)
    vad = str(vad)
    # note that this step only generates a "*.vad" file
    # this file will be of the same name but different extension (.vad)
    with open(file_path.replace(".wav", ".vad"), "w") as f:
        f.write(vad)

if len(sys.argv)!= 2:
    print("Usage: python dataset-separate.py <download_dump_dir>")
    print("The 'download_dump_dir' is the directory where the download dump files are stored.")
    print("Expecting results in the format produced by dataset-download.py")
    exit(1)

download_dump_dir = sys.argv[1]

default_output_dir = "./mdx_extra"

# note that the following are just one of many implementation possibilities

# this part assumes you keep some sort of log during download.
logs = []
for filename in os.listdir(download_dump_dir):
    if filename.endswith(".log"):
        logs.append(os.path.join(download_dump_dir, filename))
        
target_file_names = []
for log in logs:
    # here we parse the logs you kept during download.
    with open(log, "r") as f:
        lines = f.readlines()
        # we assume the first line of your download function is a target file name.
        target_file_names.append(lines[0].strip() + ".flac")

folder_names = os.listdir(default_output_dir)
for name in target_file_names:
    # if appears already in the output folder, skip
    for folder_name in folder_names:
        if folder_name == name.replace(".flac", ""):
            print(name + " appears already in the output folder, skipping...")
            target_file_names.remove(name)

for target_file_name in tqdm(target_file_names, desc="Separate Vocals"):
    # we call demucs to separate the vocals.
    os.system("demucs --two-stems=vocals -n mdx_extra \"" + os.path.join(download_dump_dir, target_file_name) + "\"" )

    # after then, we run vocals.wav through the vad defined as above.
    run_vad(os.path.join(default_output_dir, target_file_name.replace(".flac", "/vocals.wav")))
dataset/separate.py
yongyizang/SingFake
from pydub import AudioSegment
import os, shutil
from tqdm import tqdm
from concurrent.futures import ProcessPoolExecutor
import subprocess

# This is the script used for simulating T03 (codec set) from T02.

def simulate_codec(input_path, temp_path, output_path, format, bitrate="128k", codec=None):
    # Convert the input file to the temporary path using ffmpeg
    if codec:
        subprocess.run(['ffmpeg', '-i', input_path, '-acodec', codec, '-b:a', bitrate, '-y', temp_path])
    else:
        subprocess.run(['ffmpeg', '-i', input_path, '-b:a', bitrate, '-y', temp_path])
    # Read the temporary file and convert it back to flac
    subprocess.run(['ffmpeg', '-i', temp_path, '-acodec', 'flac', '-y', output_path])

    # Remove the temporary file
    os.remove(temp_path)

    
def format_suffix(format):
    if format == "mp3":
        return "mp3"
    elif format == "adts":
        return "aac"
    elif format == "ogg":
        return "ogg"
    elif format == "opus":
        return "opus"
    else:
        raise ValueError(f"Invalid format {format}")
    
def format_codec(format):
    if format == "mp3":
        return "libmp3lame"
    elif format == "ogg":
        return "libvorbis"
    elif format == "adts":
        return None
    else:
        # raise ValueError(f"Invalid format {format}")
        return None

def worker(file_tuple):
    src_file_path, dest_file_path, format, bitrate, src_file_name = file_tuple
    temp_path = os.path.join(dest_file_path.split(".flac")[0] + f".{format_suffix(format)}")
    try:
        simulate_codec(src_file_path, temp_path, dest_file_path, format=format, bitrate=bitrate, codec=format_codec(format))
    except Exception as e:
        print(f"Failed to simulate {src_file_path}")
        print(e)

def process_audio_files(src_folder, dest_folder, format="mp3", bitrate="128k", max_workers=None):
    file_list = []
    for subdir, _, files in os.walk(src_folder):
        for file in files:
            if file.endswith('.flac'):
                src_file_path = os.path.join(subdir, file)
                
                # Create the corresponding output directory
                rel_path = os.path.relpath(subdir, src_folder)
                dest_subdir = os.path.join(dest_folder, rel_path)
                os.makedirs(dest_subdir, exist_ok=True)
                
                dest_file_path = os.path.join(dest_subdir, file)
                
                file_list.append((src_file_path, dest_file_path, format, bitrate, file))

    with ProcessPoolExecutor(max_workers=max_workers) as executor:
        list(tqdm(executor.map(worker, file_list), total=len(file_list), desc=f"Processing {src_folder}"))

audio_formats = None # write audio formats you want
assert audio_formats is not None, "You must specify audio codec formats!"
# example: [["ogg", "64k"], ["opus", "64k"]]

for audio_format_tuple in audio_formats:
    audio_format = audio_format_tuple[0]
    bitrate = audio_format_tuple[1]
    print("Processing " + audio_format + "...")
    src_folder = "/home/yongyi/split_0831/test"
    dest_folder = "/home/yongyi/split_0831/codec_test/"
    dest_folder = dest_folder + audio_format + "_" + bitrate + "/"
    process_audio_files(src_folder, dest_folder, format=audio_format, bitrate=bitrate, max_workers=8)

dataset/simulate_codec.py
yongyizang/SingFake
import os, sys, librosa
import numpy as np
import soundfile as sf
from tqdm import tqdm


# customize these three paths

source_folder = "./mdx_extra" # where the separated files are. This script assumes you have ran separate.py before, and *.vad are under the same folder with file name adhereing to the same logic as written in separate.py.

logs_folder = "./logs" # It needs the *.log files (similar to separate.py)

dump_folder = "./split_dump" # output

output_sr = 16000 # sample rate for output of all files.

vad_files = []

for root, dirs, files in os.walk(source_folder):
    for file in files:
        if file.endswith(".vad"):
            vad_files.append(os.path.join(root, file))

for vad_file in tqdm(vad_files, desc="Rendering Splits"):
    vocal_only = vad_file.replace('.vad', '.wav')
    corresponding_log = os.path.join(logs_folder, vad_file.split("/")[-2].split("_")[0] + ".log")
    # here, it uses the logs file during download to determine if the file is spoof or bonafide.
    spoof = 0
    with open(corresponding_log, "r") as f:
        curr_lines = f.readlines()
        if curr_lines[4].strip() == "spoof":
            spoof = 1
        # this tag is directly written in file name!
        # dataloaders in models/(model) also adhere to this logic.
        mixture_file_name = curr_lines[0].strip() + ".flac"
    mixture_file = os.path.join(logs_folder, mixture_file_name)
    
    # open both
    vocal, vocal_sr = librosa.load(vocal_only, sr=output_sr, mono=False)
    mixture, mixture_sr = librosa.load(mixture_file, sr=output_sr, mono=False)
    
    isMono = False
    
    if len(vocal.shape) == 1:
        isMono = True
    
    if len(mixture.shape) == 1:
        isMono = True
    
    index = 0
    
    with open(vad_file, "r") as f:
        lines = f.readlines()
        for line in lines:
            line = line.strip()
            if line == "":
                continue
            else:
                try:
                    items = line.split("]")[0].split("[")[1].split("-->")
                    start_times = items[0].split(":")
                    end_times = items[1].split(":")
                    
                    start_time = float(start_times[2]) + float(start_times[1]) * 60.0 + float(start_times[0]) * 3600.0
                    end_time = float(end_times[2]) + float(end_times[1]) * 60.0 + float(end_times[0]) * 3600.0
                    if isMono:
                        vocal_curr_seg = vocal[int(start_time * output_sr):int(end_time * output_sr)]
                        mixture_curr_seg = mixture[int(start_time * output_sr):int(end_time * output_sr)]
                        vocal_curr_seg = np.expand_dims(vocal_curr_seg, axis=0)
                        mixture_curr_seg = np.expand_dims(mixture_curr_seg, axis=0)
                        
                        # transpose
                        vocal_curr_seg = np.transpose(vocal_curr_seg)
                        mixture_curr_seg = np.transpose(mixture_curr_seg)
                    else:
                        vocal_curr_seg = vocal[:, int(start_time * output_sr):int(end_time * output_sr)]
                        mixture_curr_seg = mixture[:, int(start_time * output_sr):int(end_time * output_sr)]
                        
                        # transpose
                        vocal_curr_seg = np.transpose(vocal_curr_seg)
                        mixture_curr_seg = np.transpose(mixture_curr_seg)
                    
                    file_name = str(spoof) + "_" + vad_file.split("/")[-2].split("_")[0] + "_" + str(index) + ".flac"
                    
                    # save vocal
                    sf.write(os.path.join(dump_folder, "vocals", file_name), vocal_curr_seg, vocal_sr, subtype="PCM_16")
                    
                    # save mixture
                    sf.write(os.path.join(dump_folder, "mixtures", file_name), mixture_curr_seg, mixture_sr, subtype="PCM_16")
                    
                    index += 1
                except Exception as e:
                    print(e)
                    continue
dataset/split.py
yongyizang/SingFake
"""
AASIST
Copyright (c) 2021-present NAVER Corp.
MIT license
"""

import random
from typing import Union

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor


class GraphAttentionLayer(nn.Module):
    def __init__(self, in_dim, out_dim, **kwargs):
        super().__init__()

        # attention map
        self.att_proj = nn.Linear(in_dim, out_dim)
        self.att_weight = self._init_new_params(out_dim, 1)

        # project
        self.proj_with_att = nn.Linear(in_dim, out_dim)
        self.proj_without_att = nn.Linear(in_dim, out_dim)

        # batch norm
        self.bn = nn.BatchNorm1d(out_dim)

        # dropout for inputs
        self.input_drop = nn.Dropout(p=0.2)

        # activate
        self.act = nn.SELU(inplace=True)

        # temperature
        self.temp = 1.
        if "temperature" in kwargs:
            self.temp = kwargs["temperature"]

    def forward(self, x):
        '''
        x   :(#bs, #node, #dim)
        '''
        # apply input dropout
        x = self.input_drop(x)

        # derive attention map
        att_map = self._derive_att_map(x)

        # projection
        x = self._project(x, att_map)

        # apply batch norm
        x = self._apply_BN(x)
        x = self.act(x)
        return x

    def _pairwise_mul_nodes(self, x):
        '''
        Calculates pairwise multiplication of nodes.
        - for attention map
        x           :(#bs, #node, #dim)
        out_shape   :(#bs, #node, #node, #dim)
        '''

        nb_nodes = x.size(1)
        x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1)
        x_mirror = x.transpose(1, 2)

        return x * x_mirror

    def _derive_att_map(self, x):
        '''
        x           :(#bs, #node, #dim)
        out_shape   :(#bs, #node, #node, 1)
        '''
        att_map = self._pairwise_mul_nodes(x)
        # size: (#bs, #node, #node, #dim_out)
        att_map = torch.tanh(self.att_proj(att_map))
        # size: (#bs, #node, #node, 1)
        att_map = torch.matmul(att_map, self.att_weight)

        # apply temperature
        att_map = att_map / self.temp

        att_map = F.softmax(att_map, dim=-2)

        return att_map

    def _project(self, x, att_map):
        x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x))
        x2 = self.proj_without_att(x)

        return x1 + x2

    def _apply_BN(self, x):
        org_size = x.size()
        x = x.view(-1, org_size[-1])
        x = self.bn(x)
        x = x.view(org_size)

        return x

    def _init_new_params(self, *size):
        out = nn.Parameter(torch.FloatTensor(*size))
        nn.init.xavier_normal_(out)
        return out


class HtrgGraphAttentionLayer(nn.Module):
    def __init__(self, in_dim, out_dim, **kwargs):
        super().__init__()

        self.proj_type1 = nn.Linear(in_dim, in_dim)
        self.proj_type2 = nn.Linear(in_dim, in_dim)

        # attention map
        self.att_proj = nn.Linear(in_dim, out_dim)
        self.att_projM = nn.Linear(in_dim, out_dim)

        self.att_weight11 = self._init_new_params(out_dim, 1)
        self.att_weight22 = self._init_new_params(out_dim, 1)
        self.att_weight12 = self._init_new_params(out_dim, 1)
        self.att_weightM = self._init_new_params(out_dim, 1)

        # project
        self.proj_with_att = nn.Linear(in_dim, out_dim)
        self.proj_without_att = nn.Linear(in_dim, out_dim)

        self.proj_with_attM = nn.Linear(in_dim, out_dim)
        self.proj_without_attM = nn.Linear(in_dim, out_dim)

        # batch norm
        self.bn = nn.BatchNorm1d(out_dim)

        # dropout for inputs
        self.input_drop = nn.Dropout(p=0.2)

        # activate
        self.act = nn.SELU(inplace=True)

        # temperature
        self.temp = 1.
        if "temperature" in kwargs:
            self.temp = kwargs["temperature"]

    def forward(self, x1, x2, master=None):
        '''
        x1  :(#bs, #node, #dim)
        x2  :(#bs, #node, #dim)
        '''
        num_type1 = x1.size(1)
        num_type2 = x2.size(1)

        x1 = self.proj_type1(x1)
        x2 = self.proj_type2(x2)

        x = torch.cat([x1, x2], dim=1)

        if master is None:
            master = torch.mean(x, dim=1, keepdim=True)

        # apply input dropout
        x = self.input_drop(x)

        # derive attention map
        att_map = self._derive_att_map(x, num_type1, num_type2)

        # directional edge for master node
        master = self._update_master(x, master)

        # projection
        x = self._project(x, att_map)

        # apply batch norm
        x = self._apply_BN(x)
        x = self.act(x)

        x1 = x.narrow(1, 0, num_type1)
        x2 = x.narrow(1, num_type1, num_type2)

        return x1, x2, master

    def _update_master(self, x, master):

        att_map = self._derive_att_map_master(x, master)
        master = self._project_master(x, master, att_map)

        return master

    def _pairwise_mul_nodes(self, x):
        '''
        Calculates pairwise multiplication of nodes.
        - for attention map
        x           :(#bs, #node, #dim)
        out_shape   :(#bs, #node, #node, #dim)
        '''

        nb_nodes = x.size(1)
        x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1)
        x_mirror = x.transpose(1, 2)

        return x * x_mirror

    def _derive_att_map_master(self, x, master):
        '''
        x           :(#bs, #node, #dim)
        out_shape   :(#bs, #node, #node, 1)
        '''
        att_map = x * master
        att_map = torch.tanh(self.att_projM(att_map))

        att_map = torch.matmul(att_map, self.att_weightM)

        # apply temperature
        att_map = att_map / self.temp

        att_map = F.softmax(att_map, dim=-2)

        return att_map

    def _derive_att_map(self, x, num_type1, num_type2):
        '''
        x           :(#bs, #node, #dim)
        out_shape   :(#bs, #node, #node, 1)
        '''
        att_map = self._pairwise_mul_nodes(x)
        # size: (#bs, #node, #node, #dim_out)
        att_map = torch.tanh(self.att_proj(att_map))
        # size: (#bs, #node, #node, 1)

        att_board = torch.zeros_like(att_map[:, :, :, 0]).unsqueeze(-1)

        att_board[:, :num_type1, :num_type1, :] = torch.matmul(
            att_map[:, :num_type1, :num_type1, :], self.att_weight11)
        att_board[:, num_type1:, num_type1:, :] = torch.matmul(
            att_map[:, num_type1:, num_type1:, :], self.att_weight22)
        att_board[:, :num_type1, num_type1:, :] = torch.matmul(
            att_map[:, :num_type1, num_type1:, :], self.att_weight12)
        att_board[:, num_type1:, :num_type1, :] = torch.matmul(
            att_map[:, num_type1:, :num_type1, :], self.att_weight12)

        att_map = att_board

        # att_map = torch.matmul(att_map, self.att_weight12)

        # apply temperature
        att_map = att_map / self.temp

        att_map = F.softmax(att_map, dim=-2)

        return att_map

    def _project(self, x, att_map):
        x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x))
        x2 = self.proj_without_att(x)

        return x1 + x2

    def _project_master(self, x, master, att_map):

        x1 = self.proj_with_attM(torch.matmul(
            att_map.squeeze(-1).unsqueeze(1), x))
        x2 = self.proj_without_attM(master)

        return x1 + x2

    def _apply_BN(self, x):
        org_size = x.size()
        x = x.view(-1, org_size[-1])
        x = self.bn(x)
        x = x.view(org_size)

        return x

    def _init_new_params(self, *size):
        out = nn.Parameter(torch.FloatTensor(*size))
        nn.init.xavier_normal_(out)
        return out


class GraphPool(nn.Module):
    def __init__(self, k: float, in_dim: int, p: Union[float, int]):
        super().__init__()
        self.k = k
        self.sigmoid = nn.Sigmoid()
        self.proj = nn.Linear(in_dim, 1)
        self.drop = nn.Dropout(p=p) if p > 0 else nn.Identity()
        self.in_dim = in_dim

    def forward(self, h):
        Z = self.drop(h)
        weights = self.proj(Z)
        scores = self.sigmoid(weights)
        new_h = self.top_k_graph(scores, h, self.k)

        return new_h

    def top_k_graph(self, scores, h, k):
        """
        args
        =====
        scores: attention-based weights (#bs, #node, 1)
        h: graph data (#bs, #node, #dim)
        k: ratio of remaining nodes, (float)

        returns
        =====
        h: graph pool applied data (#bs, #node', #dim)
        """
        _, n_nodes, n_feat = h.size()
        n_nodes = max(int(n_nodes * k), 1)
        _, idx = torch.topk(scores, n_nodes, dim=1)
        idx = idx.expand(-1, -1, n_feat)

        h = h * scores
        h = torch.gather(h, 1, idx)

        return h


class CONV(nn.Module):
    @staticmethod
    def to_mel(hz):
        return 2595 * np.log10(1 + hz / 700)

    @staticmethod
    def to_hz(mel):
        return 700 * (10**(mel / 2595) - 1)

    def __init__(self,
                 out_channels,
                 kernel_size,
                 sample_rate=16000,
                 in_channels=1,
                 stride=1,
                 padding=0,
                 dilation=1,
                 bias=False,
                 groups=1,
                 mask=False):
        super().__init__()
        if in_channels != 1:

            msg = "SincConv only support one input channel (here, in_channels = {%i})" % (
                in_channels)
            raise ValueError(msg)
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.sample_rate = sample_rate

        # Forcing the filters to be odd (i.e, perfectly symmetrics)
        if kernel_size % 2 == 0:
            self.kernel_size = self.kernel_size + 1
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.mask = mask
        if bias:
            raise ValueError('SincConv does not support bias.')
        if groups > 1:
            raise ValueError('SincConv does not support groups.')

        NFFT = 512
        f = int(self.sample_rate / 2) * np.linspace(0, 1, int(NFFT / 2) + 1)
        fmel = self.to_mel(f)
        fmelmax = np.max(fmel)
        fmelmin = np.min(fmel)
        filbandwidthsmel = np.linspace(fmelmin, fmelmax, self.out_channels + 1)
        filbandwidthsf = self.to_hz(filbandwidthsmel)

        self.mel = filbandwidthsf
        self.hsupp = torch.arange(-(self.kernel_size - 1) / 2,
                                  (self.kernel_size - 1) / 2 + 1)
        self.band_pass = torch.zeros(self.out_channels, self.kernel_size)
        for i in range(len(self.mel) - 1):
            fmin = self.mel[i]
            fmax = self.mel[i + 1]
            hHigh = (2*fmax/self.sample_rate) * \
                np.sinc(2*fmax*self.hsupp/self.sample_rate)
            hLow = (2*fmin/self.sample_rate) * \
                np.sinc(2*fmin*self.hsupp/self.sample_rate)
            hideal = hHigh - hLow

            self.band_pass[i, :] = Tensor(np.hamming(
                self.kernel_size)) * Tensor(hideal)

    def forward(self, x, mask=False):
        band_pass_filter = self.band_pass.clone().to(x.device)
        if mask:
            A = np.random.uniform(0, 20)
            A = int(A)
            A0 = random.randint(0, band_pass_filter.shape[0] - A)
            band_pass_filter[A0:A0 + A, :] = 0
        else:
            band_pass_filter = band_pass_filter

        self.filters = (band_pass_filter).view(self.out_channels, 1,
                                               self.kernel_size)

        return F.conv1d(x,
                        self.filters,
                        stride=self.stride,
                        padding=self.padding,
                        dilation=self.dilation,
                        bias=None,
                        groups=1)


class Residual_block(nn.Module):
    def __init__(self, nb_filts, first=False):
        super().__init__()
        self.first = first

        if not self.first:
            self.bn1 = nn.BatchNorm2d(num_features=nb_filts[0])
        self.conv1 = nn.Conv2d(in_channels=nb_filts[0],
                               out_channels=nb_filts[1],
                               kernel_size=(2, 3),
                               padding=(1, 1),
                               stride=1)
        self.selu = nn.SELU(inplace=True)

        self.bn2 = nn.BatchNorm2d(num_features=nb_filts[1])
        self.conv2 = nn.Conv2d(in_channels=nb_filts[1],
                               out_channels=nb_filts[1],
                               kernel_size=(2, 3),
                               padding=(0, 1),
                               stride=1)

        if nb_filts[0] != nb_filts[1]:
            self.downsample = True
            self.conv_downsample = nn.Conv2d(in_channels=nb_filts[0],
                                             out_channels=nb_filts[1],
                                             padding=(0, 1),
                                             kernel_size=(1, 3),
                                             stride=1)

        else:
            self.downsample = False
        self.mp = nn.MaxPool2d((1, 3))  # self.mp = nn.MaxPool2d((1,4))

    def forward(self, x):
        identity = x
        if not self.first:
            out = self.bn1(x)
            out = self.selu(out)
        else:
            out = x
        out = self.conv1(x)

        # print('out',out.shape)
        out = self.bn2(out)
        out = self.selu(out)
        # print('out',out.shape)
        out = self.conv2(out)
        #print('conv2 out',out.shape)
        if self.downsample:
            identity = self.conv_downsample(identity)

        out += identity
        out = self.mp(out)
        return out


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

        self.d_args = d_args
        filts = d_args["filts"]
        gat_dims = d_args["gat_dims"]
        pool_ratios = d_args["pool_ratios"]
        temperatures = d_args["temperatures"]

        self.conv_time = CONV(out_channels=filts[0],
                              kernel_size=d_args["first_conv"],
                              in_channels=1)
        self.first_bn = nn.BatchNorm2d(num_features=1)

        self.drop = nn.Dropout(0.5, inplace=True)
        self.drop_way = nn.Dropout(0.2, inplace=True)
        self.selu = nn.SELU(inplace=True)

        self.encoder = nn.Sequential(
            nn.Sequential(Residual_block(nb_filts=filts[1], first=True)),
            nn.Sequential(Residual_block(nb_filts=filts[2])),
            nn.Sequential(Residual_block(nb_filts=filts[3])),
            nn.Sequential(Residual_block(nb_filts=filts[4])),
            nn.Sequential(Residual_block(nb_filts=filts[4])),
            nn.Sequential(Residual_block(nb_filts=filts[4])))

        self.pos_S = nn.Parameter(torch.randn(1, 23, filts[-1][-1]))
        self.master1 = nn.Parameter(torch.randn(1, 1, gat_dims[0]))
        self.master2 = nn.Parameter(torch.randn(1, 1, gat_dims[0]))

        self.GAT_layer_S = GraphAttentionLayer(filts[-1][-1],
                                               gat_dims[0],
                                               temperature=temperatures[0])
        self.GAT_layer_T = GraphAttentionLayer(filts[-1][-1],
                                               gat_dims[0],
                                               temperature=temperatures[1])

        self.HtrgGAT_layer_ST11 = HtrgGraphAttentionLayer(
            gat_dims[0], gat_dims[1], temperature=temperatures[2])
        self.HtrgGAT_layer_ST12 = HtrgGraphAttentionLayer(
            gat_dims[1], gat_dims[1], temperature=temperatures[2])

        self.HtrgGAT_layer_ST21 = HtrgGraphAttentionLayer(
            gat_dims[0], gat_dims[1], temperature=temperatures[2])

        self.HtrgGAT_layer_ST22 = HtrgGraphAttentionLayer(
            gat_dims[1], gat_dims[1], temperature=temperatures[2])

        self.pool_S = GraphPool(pool_ratios[0], gat_dims[0], 0.3)
        self.pool_T = GraphPool(pool_ratios[1], gat_dims[0], 0.3)
        self.pool_hS1 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)
        self.pool_hT1 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)

        self.pool_hS2 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)
        self.pool_hT2 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)

        self.out_layer = nn.Linear(5 * gat_dims[1], 2)

    def forward(self, x, Freq_aug=False):

        x = x.unsqueeze(1)
        x = self.conv_time(x, mask=Freq_aug)
        x = x.unsqueeze(dim=1)
        x = F.max_pool2d(torch.abs(x), (3, 3))
        x = self.first_bn(x)
        x = self.selu(x)

        # get embeddings using encoder
        # (#bs, #filt, #spec, #seq)
        e = self.encoder(x)

        # spectral GAT (GAT-S)
        e_S, _ = torch.max(torch.abs(e), dim=3)  # max along time
        e_S = e_S.transpose(1, 2) + self.pos_S

        gat_S = self.GAT_layer_S(e_S)
        out_S = self.pool_S(gat_S)  # (#bs, #node, #dim)

        # temporal GAT (GAT-T)
        e_T, _ = torch.max(torch.abs(e), dim=2)  # max along freq
        e_T = e_T.transpose(1, 2)

        gat_T = self.GAT_layer_T(e_T)
        out_T = self.pool_T(gat_T)

        # learnable master node
        master1 = self.master1.expand(x.size(0), -1, -1)
        master2 = self.master2.expand(x.size(0), -1, -1)

        # inference 1
        out_T1, out_S1, master1 = self.HtrgGAT_layer_ST11(
            out_T, out_S, master=self.master1)

        out_S1 = self.pool_hS1(out_S1)
        out_T1 = self.pool_hT1(out_T1)

        out_T_aug, out_S_aug, master_aug = self.HtrgGAT_layer_ST12(
            out_T1, out_S1, master=master1)
        out_T1 = out_T1 + out_T_aug
        out_S1 = out_S1 + out_S_aug
        master1 = master1 + master_aug

        # inference 2
        out_T2, out_S2, master2 = self.HtrgGAT_layer_ST21(
            out_T, out_S, master=self.master2)
        out_S2 = self.pool_hS2(out_S2)
        out_T2 = self.pool_hT2(out_T2)

        out_T_aug, out_S_aug, master_aug = self.HtrgGAT_layer_ST22(
            out_T2, out_S2, master=master2)
        out_T2 = out_T2 + out_T_aug
        out_S2 = out_S2 + out_S_aug
        master2 = master2 + master_aug

        out_T1 = self.drop_way(out_T1)
        out_T2 = self.drop_way(out_T2)
        out_S1 = self.drop_way(out_S1)
        out_S2 = self.drop_way(out_S2)
        master1 = self.drop_way(master1)
        master2 = self.drop_way(master2)

        out_T = torch.max(out_T1, out_T2)
        out_S = torch.max(out_S1, out_S2)
        master = torch.max(master1, master2)

        T_max, _ = torch.max(torch.abs(out_T), dim=1)
        T_avg = torch.mean(out_T, dim=1)

        S_max, _ = torch.max(torch.abs(out_S), dim=1)
        S_avg = torch.mean(out_S, dim=1)

        last_hidden = torch.cat(
            [T_max, T_avg, S_max, S_avg, master.squeeze(1)], dim=1)

        last_hidden = self.drop(last_hidden)
        output = self.out_layer(last_hidden)

        return output

models/AASIST/AASIST.py
yongyizang/SingFake
import numpy as np
import soundfile as sf
import torch, os
from torch import Tensor
from torch.utils.data import Dataset
import librosa

___author__ = "Hemlata Tak, Jee-weon Jung"
__email__ = "tak@eurecom.fr, jeeweon.jung@navercorp.com"


def genSpoof_list(dir_meta, is_train=False, is_eval=False):
    d_meta = {}
    file_list = []
    with open(dir_meta, "r") as f:
        l_meta = f.readlines()

    if is_train:
        for line in l_meta:
            _, key, _, _, label = line.strip().split(" ")
            file_list.append(key)
            d_meta[key] = 1 if label == "bonafide" else 0
        return d_meta, file_list

    elif is_eval:
        for line in l_meta:
            _, key, _, _, _ = line.strip().split(" ")
            #key = line.strip()
            file_list.append(key)
        return file_list
    else:
        for line in l_meta:
            _, key, _, _, label = line.strip().split(" ")
            file_list.append(key)
            d_meta[key] = 1 if label == "bonafide" else 0
        return d_meta, file_list


def pad(x, max_len=64600):
    x_len = x.shape[0]
    if x_len >= max_len:
        return x[:max_len]
    # need to pad
    num_repeats = int(max_len / x_len) + 1
    padded_x = np.tile(x, (1, num_repeats))[:, :max_len][0]
    return padded_x


def pad_random(x: np.ndarray, max_len: int = 64600):
    x_len = x.shape[0]
    # if duration is already long enough
    if x_len >= max_len:
        stt = np.random.randint(x_len - max_len)
        return x[stt:stt + max_len]

    # if too short
    num_repeats = int(max_len / x_len) + 1
    padded_x = np.tile(x, (num_repeats))[:max_len]
    return padded_x


class Dataset_ASVspoof2019_train(Dataset):
    def __init__(self, list_IDs, labels, base_dir):
        """self.list_IDs	: list of strings (each string: utt key),
           self.labels      : dictionary (key: utt key, value: label integer)"""
        self.list_IDs = list_IDs
        self.labels = labels
        self.base_dir = base_dir
        self.cut = 64600  # take ~4 sec audio (64600 samples)

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

    def __getitem__(self, index):
        key = self.list_IDs[index]
        X, _ = sf.read(str(self.base_dir / f"flac/{key}.flac"))
        X_pad = pad_random(X, self.cut)
        x_inp = Tensor(X_pad)
        y = self.labels[key]
        return x_inp, y


class Dataset_ASVspoof2019_devNeval(Dataset):
    def __init__(self, list_IDs, base_dir):
        """self.list_IDs	: list of strings (each string: utt key),
        """
        self.list_IDs = list_IDs
        self.base_dir = base_dir
        self.cut = 64600  # take ~4 sec audio (64600 samples)

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

    def __getitem__(self, index):
        key = self.list_IDs[index]
        X, _ = sf.read(str(self.base_dir / f"flac/{key}.flac"))
        X_pad = pad(X, self.cut)
        x_inp = Tensor(X_pad)
        return x_inp, key

class Dataset_SingFake(Dataset):
    def __init__(self, base_dir, is_mixture=False, target_sr=16000):
        """
        base_dir should contain mixtures/ and vocals/ folders
        """
        self.base_dir = base_dir
        self.is_mixture = is_mixture
        self.target_sr = target_sr
        self.cut = 64600  # take ~4 sec audio (64600 samples)
        
        # get file list
        self.file_list = []
        if self.is_mixture:
            self.target_path = os.path.join(self.base_dir, "mixtures")
        else:
            self.target_path = os.path.join(self.base_dir, "vocals")
            
        print(self.target_path)
        
        assert os.path.exists(self.target_path), f"{self.target_path} does not exist!"
        
        for file in os.listdir(self.target_path):
            if file.endswith(".flac"):
                self.file_list.append(file[:-5])
    
    def __len__(self):
        return len(self.file_list)

    def __getitem__(self, index):
        key = self.file_list[index]
        file_path = os.path.join(self.target_path, key + ".flac")
        # X, _ = sf.read(file_path, samplerate=self.target_sr)
        try:
            X, _ = librosa.load(file_path, sr=self.target_sr, mono=False)
        except:
            return self.__getitem__(np.random.randint(len(self.file_list)))
        if X.shape[0] > 1:
            # if not mono, take random channel
            channel_id = np.random.randint(X.shape[0])
            X = X[channel_id]
            # X = np.expand_dims(X, axis=-1)
        X_pad = pad_random(X, self.cut)
        X_pad = X_pad / np.max(np.abs(X_pad))
        x_inp = Tensor(X_pad)
        y = int(key.split("_")[0])
        return x_inp, y
models/AASIST/data_utils.py
yongyizang/SingFake
"""
Main script that trains, validates, and evaluates
various models including AASIST.

AASIST
Copyright (c) 2021-present NAVER Corp.
MIT license
"""
import argparse
import json
import os
import sys
import warnings
from importlib import import_module
from pathlib import Path
from shutil import copy
from typing import Dict, List, Union
from tqdm import tqdm
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchcontrib.optim import SWA
from data_utils import (Dataset_ASVspoof2019_train,
                        Dataset_ASVspoof2019_devNeval, genSpoof_list, Dataset_SingFake)
from evaluation import compute_eer
from utils import create_optimizer, seed_worker, set_seed, str_to_bool
from models.wav2vecAASIST import Wav2Vec2Model
warnings.filterwarnings("ignore", category=FutureWarning)

def main(args: argparse.Namespace) -> None:
    """
    Main function.
    Trains, validates, and evaluates the ASVspoof detection model.
    """
    # load experiment configurations
    with open(args.config, "r") as f_json:
        config = json.loads(f_json.read())
    model_config = config["model_config"]
    optim_config = config["optim_config"]
    optim_config["epochs"] = config["num_epochs"]
    track = config["track"]
    assert track in ["LA", "PA", "DF"], "Invalid track given"
    if "eval_all_best" not in config:
        config["eval_all_best"] = "True"
    if "freq_aug" not in config:
        config["freq_aug"] = "False"

    # make experiment reproducible
    set_seed(args.seed, config)

    # define database related paths
    output_dir = Path(args.output_dir)
    prefix_2019 = "ASVspoof2019.{}".format(track)
    database_path = Path(config["database_path"])
    dev_trial_path = (database_path /
                      "ASVspoof2019_{}_cm_protocols/{}.cm.dev.trl.txt".format(
                          track, prefix_2019))
    eval_trial_path = (
        database_path /
        "ASVspoof2019_{}_cm_protocols/{}.cm.eval.trl.txt".format(
            track, prefix_2019))

    # define model related paths
    model_tag = "{}_{}_ep{}_bs{}".format(
        track,
        os.path.splitext(os.path.basename(args.config))[0],
        config["num_epochs"], config["batch_size"])
    if args.comment:
        model_tag = model_tag + "_{}".format(args.comment)
    model_tag = output_dir / model_tag
    model_save_path = model_tag / "weights"
    eval_score_path = model_tag / config["eval_output"]
    writer = SummaryWriter(model_tag)
    os.makedirs(model_save_path, exist_ok=True)
    copy(args.config, model_tag / "config.conf")

    # set device
    gpu_id = args.gpu
    device = torch.device("cuda:{}".format(gpu_id) if torch.cuda.is_available() else "cpu")
    # device = "cuda" if torch.cuda.is_available() else "cpu"
    print("Device: {}".format(device))
    if device == "cpu":
        raise ValueError("GPU not detected!")

    # define model architecture
    # model = get_model(model_config, device)
    model = get_wav2vec2_model(model_config, device).to(device)

    # define dataloaders
    trn_loader, dev_loader, eval_loader, additional_loader, persian_loader, mp3_loader, ogg_loader, aac_loader, opus_loader = get_singfake_loaders(args.seed, config)

    # evaluates pretrained model and exit script
    if args.eval:
        model.load_state_dict(
            torch.load(config["model_path"], map_location=device))
        print("Model loaded : {}".format(config["model_path"]))
        print("Evaluating on SingFake Eval Set...")
        train_eer = evaluate(trn_loader, model, device)
        test_eer = evaluate(eval_loader, model, device)
        additional_set_eer = evaluate(additional_loader, model, device)
        persian_eer = evaluate(persian_loader, model, device)
        mp3_eer = evaluate(mp3_loader, model, device)
        ogg_eer = evaluate(ogg_loader, model, device)
        aac_eer = evaluate(aac_loader, model, device)
        opus_eer = evaluate(opus_loader, model, device)
        average_codec_eer = (mp3_eer + aac_eer + ogg_eer + opus_eer) / 4.
        print("Done. train_eer: {:.2f} %, test_set_eer: {:.2f} %, additional_test_eer: {:.2f} %, persian_eer: {:.2f} %".format(train_eer * 100, test_eer * 100, additional_set_eer * 100, persian_eer * 100))
        print("Codec testing: Average EER: {:.2f} %, MP3 EER: {:.2f} %, OGG EER: {:.2f} %, AAC EER: {:.2f} %, OPUS EER: {:.2f} %".format(average_codec_eer * 100, mp3_eer * 100, ogg_eer * 100, aac_eer * 100, opus_eer * 100))
        sys.exit(0)
        
    # get optimizer and scheduler
    optim_config["steps_per_epoch"] = len(trn_loader)
    optimizer, scheduler = create_optimizer(model.parameters(), optim_config)
    optimizer_swa = SWA(optimizer)

    best_dev_eer = 1.
    best_eval_eer = 100.
    n_swa_update = 0  # number of snapshots of model to use in SWA
    f_log = open(model_tag / "metric_log.txt", "a")
    f_log.write("=" * 5 + "\n")

    # make directory for metric logging
    metric_path = model_tag / "metrics"
    os.makedirs(metric_path, exist_ok=True)

    # Training
    for epoch in range(config["num_epochs"]):
        print("Start training epoch{:03d}".format(epoch))
        running_loss = train_epoch(trn_loader, model, optimizer, device,
                                   scheduler, config)
        
        train_eer = evaluate(trn_loader, model, device)
        dev_eer = evaluate(dev_loader, model, device)
        additional_eer = evaluate(additional_loader, model, device)
        
        print("DONE.\nLoss:{:.5f}, train_eer: {:2f} %, dev_eer: {:.2f} %, additional_test_eer: {:.2f} %".format(
            running_loss, train_eer * 100, dev_eer * 100, additional_eer * 100))
        writer.add_scalar("loss", running_loss, epoch)
        writer.add_scalar("train_eer", train_eer, epoch)
        writer.add_scalar("dev_eer", dev_eer, epoch)
        writer.add_scalar("additional_eer", additional_eer, epoch)
        
        if best_dev_eer >= dev_eer:
            print("best model find at epoch", epoch)
            best_dev_eer = dev_eer
            torch.save(model.state_dict(),
                       model_save_path / "epoch_{}_{:03.3f}.pth".format(epoch, dev_eer))

            # do evaluation whenever best model is renewed
            if str_to_bool(config["eval_all_best"]):
                eval_eer = evaluate(eval_loader, model, device)
                log_text = "epoch{:03d}, ".format(epoch)
                if eval_eer < best_eval_eer:
                    log_text += "best eer, {:.4f}%".format(eval_eer)
                    best_eval_eer = eval_eer
                    torch.save(model.state_dict(), model_save_path / "best.pth")
                    
                if len(log_text) > 0:
                    print(log_text)
                    f_log.write(log_text + "\n")

            print("Saving epoch {} for swa".format(epoch))
            optimizer_swa.update_swa()
            n_swa_update += 1
        writer.add_scalar("best_dev_eer", best_dev_eer, epoch)

    print("Start final evaluation")
    epoch += 1
    if n_swa_update > 0:
        optimizer_swa.swap_swa_sgd()
        optimizer_swa.bn_update(trn_loader, model, device=device)
    eval_eer = evaluate(eval_loader, model, device)
    f_log = open(model_tag / "metric_log.txt", "a")
    f_log.write("=" * 5 + "\n")
    f_log.write("EER: {:.3f} %".format(eval_eer * 100))
    f_log.close()

    torch.save(model.state_dict(),
               model_save_path / "swa.pth")

    if eval_eer <= best_eval_eer:
        best_eval_eer = eval_eer
        torch.save(model.state_dict(), model_save_path / "best.pth")
        
    print("Exp FIN. EER: {:.3f} %".format(
        best_eval_eer * 100))


def get_model(model_config: Dict, device: torch.device):
    """Define DNN model architecture"""
    module = import_module("models.{}".format(model_config["architecture"]))
    _model = getattr(module, "Model")
    model = _model(model_config).to(device)
    nb_params = sum([param.view(-1).size()[0] for param in model.parameters()])
    print("no. model params:{}".format(nb_params))

    return model

def get_wav2vec2_model(model_config: Dict, device: torch.device):
    model = Wav2Vec2Model(model_config, device)
    return model


def get_loader(
        database_path: str,
        seed: int,
        config: dict) -> List[torch.utils.data.DataLoader]:
    """Make PyTorch DataLoaders for train / developement / evaluation"""
    track = config["track"]
    prefix_2019 = "ASVspoof2019.{}".format(track)

    trn_database_path = database_path / "ASVspoof2019_{}_train/".format(track)
    dev_database_path = database_path / "ASVspoof2019_{}_dev/".format(track)
    eval_database_path = database_path / "ASVspoof2019_{}_eval/".format(track)

    trn_list_path = (database_path /
                     "ASVspoof2019_{}_cm_protocols/{}.cm.train.trn.txt".format(
                         track, prefix_2019))
    dev_trial_path = (database_path /
                      "ASVspoof2019_{}_cm_protocols/{}.cm.dev.trl.txt".format(
                          track, prefix_2019))
    eval_trial_path = (
        database_path /
        "ASVspoof2019_{}_cm_protocols/{}.cm.eval.trl.txt".format(
            track, prefix_2019))

    d_label_trn, file_train = genSpoof_list(dir_meta=trn_list_path,
                                            is_train=True,
                                            is_eval=False)
    print("no. training files:", len(file_train))

    train_set = Dataset_ASVspoof2019_train(list_IDs=file_train,
                                           labels=d_label_trn,
                                           base_dir=trn_database_path)
    gen = torch.Generator()
    gen.manual_seed(seed)
    trn_loader = DataLoader(train_set,
                            batch_size=config["batch_size"],
                            shuffle=True,
                            drop_last=True,
                            pin_memory=True,
                            worker_init_fn=seed_worker,
                            generator=gen)

    _, file_dev = genSpoof_list(dir_meta=dev_trial_path,
                                is_train=False,
                                is_eval=False)
    print("no. validation files:", len(file_dev))

    dev_set = Dataset_ASVspoof2019_devNeval(list_IDs=file_dev,
                                            base_dir=dev_database_path)
    dev_loader = DataLoader(dev_set,
                            batch_size=config["batch_size"],
                            shuffle=False,
                            drop_last=False,
                            pin_memory=True)

    file_eval = genSpoof_list(dir_meta=eval_trial_path,
                              is_train=False,
                              is_eval=True)
    eval_set = Dataset_ASVspoof2019_devNeval(list_IDs=file_eval,
                                             base_dir=eval_database_path)
    eval_loader = DataLoader(eval_set,
                             batch_size=config["batch_size"],
                             shuffle=False,
                             drop_last=False,
                             pin_memory=True)

    return trn_loader, dev_loader, eval_loader

def get_singfake_loaders(seed: int, config: dict) -> List[torch.utils.data.DataLoader]:
    """
    Work in progress; no split is provided as of now.
    """
    base_dir = "/home/yongyi/split_0831"
    
    target_sr = float(config["target_sr"])
    vocals_only = str_to_bool(config["vocals_only"])
    
    is_mixture = not vocals_only
    
    train_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "train"), is_mixture=is_mixture, target_sr=target_sr)
    gen = torch.Generator()
    gen.manual_seed(seed)
    trn_loader = DataLoader(train_set,
                            batch_size=config["batch_size"],
                            shuffle=True,
                            drop_last=True,
                            pin_memory=True,
                            worker_init_fn=seed_worker,
                            num_workers=4,
                            generator=gen)
    
    dev_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "dev"), is_mixture=is_mixture, target_sr=target_sr)
    dev_loader = DataLoader(dev_set,
                            batch_size=config["batch_size"],
                            shuffle=False,
                            drop_last=False,
                            num_workers=4,
                            pin_memory=True)
    
    eval_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "test"), is_mixture=is_mixture, target_sr=target_sr)
    eval_loader = DataLoader(eval_set,
                             batch_size=config["batch_size"],
                             shuffle=False,
                             drop_last=False,
                             num_workers=4,
                             pin_memory=True)
    
    additional_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "additional_test"), is_mixture=is_mixture, target_sr=target_sr)
    additional_loader = DataLoader(additional_set,
                                      batch_size=config["batch_size"],
                                      shuffle=False,
                                      drop_last=False,
                                      num_workers=4,
                                      pin_memory=True)
    
    persian_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "persian_test"), is_mixture=is_mixture, target_sr=target_sr)
    persian_loader = DataLoader(persian_set,
                                batch_size=config["batch_size"],
                                      shuffle=False,
                                      drop_last=False,
                                      num_workers=4,
                                      pin_memory=True)
    
    mp3_test_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "codec_test", "mp3_128k"), is_mixture=is_mixture, target_sr=target_sr)
    mp3_loader = DataLoader(mp3_test_set,
                            batch_size=config["batch_size"],
                                      shuffle=False,
                                      drop_last=False,
                                      num_workers=4,
                                      pin_memory=True)
    ogg_test_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "codec_test", "ogg_64k"), is_mixture=is_mixture, target_sr=target_sr)
    ogg_loader = DataLoader(ogg_test_set,
                            batch_size=config["batch_size"],
                                      shuffle=False,
                                      drop_last=False,
                                      num_workers=4,
                                      pin_memory=True)
    aac_test_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "codec_test", "adts_64k"), is_mixture=is_mixture, target_sr=target_sr)
    aac_loader = DataLoader(aac_test_set,
                            batch_size=config["batch_size"],
                                      shuffle=False,
                                      drop_last=False,
                                      num_workers=4,
                                      pin_memory=True)
    opus_test_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "codec_test", "opus_64k"), is_mixture=is_mixture, target_sr=target_sr)
    opus_loader = DataLoader(opus_test_set,
                            batch_size=config["batch_size"],
                                      shuffle=False,
                                      drop_last=False,
                                      num_workers=4,
                                      pin_memory=True)
    
    return trn_loader, dev_loader, eval_loader, additional_loader, persian_loader, mp3_loader, ogg_loader, aac_loader, opus_loader
    


def evaluate(loader, model, device: torch.device):
    """
    Evaluate the model on the given loader, then return EER.
    """
    model.eval()
    # we save target (1) scores to target_scores, and non target (0) scores to nontarget_scores.
    target_scores = []
    nontarget_scores = []
    debug = False
    count = 0
    with torch.no_grad():
        for batch_x, batch_y in tqdm(loader, total=len(loader)):
            batch_x = batch_x.to(device)
            batch_out = model(batch_x)
            batch_score = (batch_out[:, 1]).data.cpu().numpy().ravel()
            batch_y = batch_y.data.cpu().numpy().ravel()
            for i in range(len(batch_y)):
                if batch_y[i] == 1:
                    target_scores.append(batch_score[i])
                else:
                    nontarget_scores.append(batch_score[i])
            count += 1
            if count == 10 and debug:
                break
    
    eer, _ = compute_eer(target_scores, nontarget_scores)
    return eer


def train_epoch(
    trn_loader: DataLoader,
    model,
    optim: Union[torch.optim.SGD, torch.optim.Adam],
    device: torch.device,
    scheduler: torch.optim.lr_scheduler,
    config: argparse.Namespace):
    """Train the model for one epoch"""
    running_loss = 0
    num_total = 0.0
    ii = 0
    model.train()

    # set objective (Loss) functions
    weight = torch.FloatTensor([0.1, 0.9]).to(device)
    criterion = nn.CrossEntropyLoss(weight=weight)
    pbar = tqdm(trn_loader, total=len(trn_loader))
    for batch_x, batch_y in pbar:
        batch_size = batch_x.size(0)
        num_total += batch_size
        ii += 1
        batch_x = batch_x.to(device)
        batch_y = batch_y.view(-1).type(torch.int64).to(device)
        batch_out = model(batch_x)
        batch_loss = criterion(batch_out, batch_y)
        running_loss += batch_loss.item() * batch_size
        pbar.set_description("loss: {:.5f}, running loss: {:.5f}".format(
            batch_loss.item(), running_loss / num_total))
        optim.zero_grad()
        batch_loss.backward()
        optim.step()

        if config["optim_config"]["scheduler"] in ["cosine", "keras_decay"]:
            scheduler.step()
        elif scheduler is None:
            pass
        else:
            raise ValueError("scheduler error, got:{}".format(scheduler))

    running_loss /= num_total
    return running_loss


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="ASVspoof detection system")
    parser.add_argument("--config",
                        dest="config",
                        type=str,
                        help="configuration file",
                        required=True)
    parser.add_argument(
        "--output_dir",
        dest="output_dir",
        type=str,
        help="output directory for results",
        default="./exp_result",
    )
    parser.add_argument("--seed",
                        type=int,
                        default=1234,
                        help="random seed (default: 1234)")
    parser.add_argument(
        "--eval",
        action="store_true",
        help="when this flag is given, evaluates given model and exit")
    parser.add_argument("--comment",
                        type=str,
                        default=None,
                        help="comment to describe the saved model")
    parser.add_argument("--eval_model_weights",
                        type=str,
                        default=None,
                        help="directory to the model weight file (can be also given in the config file)")
    parser.add_argument("--gpu",
                        type=int,
                        default=0,
                        help="gpu id to use (default: 0)")
    main(parser.parse_args())

models/AASIST/main.py
yongyizang/SingFake
"""
Utilization functions
"""

import os
import random
import sys

import numpy as np
import torch


def str_to_bool(val):
    """Convert a string representation of truth to true (1) or false (0).
    Copied from the python implementation distutils.utils.strtobool

    True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values
    are 'n', 'no', 'f', 'false', 'off', and '0'.  Raises ValueError if
    'val' is anything else.
    >>> str_to_bool('YES')
    1
    >>> str_to_bool('FALSE')
    0
    """
    val = val.lower()
    if val in ('y', 'yes', 't', 'true', 'on', '1'):
        return True
    if val in ('n', 'no', 'f', 'false', 'off', '0'):
        return False
    raise ValueError('invalid truth value {}'.format(val))


def cosine_annealing(step, total_steps, lr_max, lr_min):
    """Cosine Annealing for learning rate decay scheduler"""
    return lr_min + (lr_max -
                     lr_min) * 0.5 * (1 + np.cos(step / total_steps * np.pi))


def keras_decay(step, decay=0.0001):
    """Learning rate decay in Keras-style"""
    return 1. / (1. + decay * step)


class SGDRScheduler(torch.optim.lr_scheduler._LRScheduler):
    """SGD with restarts scheduler"""
    def __init__(self, optimizer, T0, T_mul, eta_min, last_epoch=-1):
        self.Ti = T0
        self.T_mul = T_mul
        self.eta_min = eta_min

        self.last_restart = 0

        super().__init__(optimizer, last_epoch)

    def get_lr(self):
        T_cur = self.last_epoch - self.last_restart
        if T_cur >= self.Ti:
            self.last_restart = self.last_epoch
            self.Ti = self.Ti * self.T_mul
            T_cur = 0

        return [
            self.eta_min + (base_lr - self.eta_min) *
            (1 + np.cos(np.pi * T_cur / self.Ti)) / 2
            for base_lr in self.base_lrs
        ]


def _get_optimizer(model_parameters, optim_config):
    """Defines optimizer according to the given config"""
    optimizer_name = optim_config['optimizer']

    if optimizer_name == 'sgd':
        optimizer = torch.optim.SGD(model_parameters,
                                    lr=optim_config['base_lr'],
                                    momentum=optim_config['momentum'],
                                    weight_decay=optim_config['weight_decay'],
                                    nesterov=optim_config['nesterov'])
    elif optimizer_name == 'adam':
        optimizer = torch.optim.Adam(model_parameters,
                                     lr=optim_config['base_lr'],
                                     betas=optim_config['betas'],
                                     weight_decay=optim_config['weight_decay'],
                                     amsgrad=str_to_bool(
                                         optim_config['amsgrad']))
    else:
        print('Un-known optimizer', optimizer_name)
        sys.exit()

    return optimizer


def _get_scheduler(optimizer, optim_config):
    """
    Defines learning rate scheduler according to the given config
    """
    if optim_config['scheduler'] == 'multistep':
        scheduler = torch.optim.lr_scheduler.MultiStepLR(
            optimizer,
            milestones=optim_config['milestones'],
            gamma=optim_config['lr_decay'])

    elif optim_config['scheduler'] == 'sgdr':
        scheduler = SGDRScheduler(optimizer, optim_config['T0'],
                                  optim_config['Tmult'],
                                  optim_config['lr_min'])

    elif optim_config['scheduler'] == 'cosine':
        total_steps = optim_config['epochs'] * \
            optim_config['steps_per_epoch']

        scheduler = torch.optim.lr_scheduler.LambdaLR(
            optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                total_steps,
                1,  # since lr_lambda computes multiplicative factor
                optim_config['lr_min'] / optim_config['base_lr']))

    elif optim_config['scheduler'] == 'keras_decay':
        scheduler = torch.optim.lr_scheduler.LambdaLR(
            optimizer, lr_lambda=lambda step: keras_decay(step))
    else:
        scheduler = None
    return scheduler


def create_optimizer(model_parameters, optim_config):
    """Defines an optimizer and a scheduler"""
    optimizer = _get_optimizer(model_parameters, optim_config)
    scheduler = _get_scheduler(optimizer, optim_config)
    return optimizer, scheduler


def seed_worker(worker_id):
    """
    Used in generating seed for the worker of torch.utils.data.Dataloader
    """
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)


def set_seed(seed, config = None):
    """ 
    set initial seed for reproduction
    """
    if config is None:
        raise ValueError("config should not be None")

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = str_to_bool(config["cudnn_deterministic_toggle"])
        torch.backends.cudnn.benchmark = str_to_bool(config["cudnn_benchmark_toggle"])

models/AASIST/utils.py
yongyizang/SingFake
import torch.nn as nn
import torch


class MaxFeatureMap2D(nn.Module):
    """ Max feature map (along 2D)
    MaxFeatureMap2D(max_dim=1)
    l_conv2d = MaxFeatureMap2D(1)
    data_in = torch.rand([1, 4, 5, 5])
    data_out = l_conv2d(data_in)
    Input:
    ------
    data_in: tensor of shape (batch, channel, ...)
    Output:
    -------
    data_out: tensor of shape (batch, channel//2, ...)
    Note
    ----
    By default, Max-feature-map is on channel dimension,
    and maxout is used on (channel ...)
    """

    def __init__(self):
        super(MaxFeatureMap2D, self).__init__()

    def forward(self, inputs):
        # suppose inputs (batchsize, channel, length, dim)

        # shape = list(inputs.size())
        # # print(shape)
        # shape[1] = shape[1] // 2
        # # print(shape)
        # shape.insert(1, 2)
        # # print(shape)

        lst = torch.split(inputs, inputs.shape[1]//2, dim=1)
        # print(lst[0].shape)
        inputs = torch.stack(lst, dim=1)
        # print(inputs.shape)

        # view to (batchsize, 2, channel//2, ...)
        # maximize on the 2nd dim
        m, i = inputs.max(1)
        return m
models/feat+resnet/LCNNhelper.py
yongyizang/SingFake
import torch
import torch.nn as torch_nn
import torch.nn.functional as torch_nn_func
import numpy as np
import librosa
import pickle
import sys

## Adapted from https://github.com/nii-yamagishilab/project-NN-Pytorch-scripts/blob/newfunctions/

__author__ = "Xin Wang"
__email__ = "wangxin@nii.ac.jp"
__copyright__ = "Copyright 2020, Xin Wang"

##################
## other utilities
##################
def trimf(x, params):
    """
    trimf: similar to Matlab definition
    https://www.mathworks.com/help/fuzzy/trimf.html?s_tid=srchtitle

    """
    if len(params) != 3:
        print("trimp requires params to be a list of 3 elements")
        sys.exit(1)
    a = params[0]
    b = params[1]
    c = params[2]
    if a > b or b > c:
        print("trimp(x, [a, b, c]) requires a<=b<=c")
        sys.exit(1)
    y = torch.zeros_like(x, dtype=torch.float32)
    if a < b:
        index = np.logical_and(a < x, x < b).bool()
        y[index] = (x[index] - a) / (b - a)
    if b < c:
        index = np.logical_and(b < x, x < c).bool()
        y[index] = (c - x[index]) / (c - b)
    y[x == b] = 1
    return y

def delta(x):
    """ By default
    input
    -----
    x (batch, Length, dim)

    output
    ------
    output (batch, Length, dim)

    Delta is calculated along Length
    """
    length = x.shape[1]
    x_temp = torch_nn_func.pad(x.unsqueeze(1), (0, 0, 1, 1),
                               'replicate').squeeze(1)
    output = -1 * x_temp[:, 0:length] + x_temp[:, 2:]
    return output


class LFCC(torch_nn.Module):
    """ Based on asvspoof.org baseline Matlab code.
    Difference: with_energy is added to set the first dimension as energy

    """

    def __init__(self, fl, fs, fn, sr, filter_num,
                 with_energy=False, with_emphasis=True,
                 with_delta=True):
        super(LFCC, self).__init__()
        self.fl = fl
        self.fs = fs
        self.fn = fn
        self.sr = sr
        self.filter_num = filter_num

        f = (sr / 2) * torch.linspace(0, 1, fn // 2 + 1)
        filter_bands = torch.linspace(min(f), max(f), filter_num + 2)

        filter_bank = torch.zeros([fn // 2 + 1, filter_num])
        for idx in range(filter_num):
            filter_bank[:, idx] = trimf(
                f, [filter_bands[idx],
                    filter_bands[idx + 1],
                    filter_bands[idx + 2]])
        self.lfcc_fb = torch_nn.Parameter(filter_bank, requires_grad=False)
        self.l_dct = LinearDCT(filter_num, 'dct', norm='ortho')
        self.with_energy = with_energy
        self.with_emphasis = with_emphasis
        self.with_delta = with_delta
        self.register_buffer("window", torch.hamming_window(self.fl))
        return

    def forward(self, x):
        """

        input:
        ------
         x: tensor(batch, length), where length is waveform length

        output:
        -------
         lfcc_output: tensor(batch, frame_num, dim_num)
        """
        # pre-emphasis
        if self.with_emphasis:
            x[:, 1:] = x[:, 1:] - 0.97 * x[:, 0:-1]

        # STFT

        x_stft = torch.stft(x, self.fn, self.fs, self.fl,
                            window=self.window,
                            onesided=True, pad_mode="constant",
                            return_complex=False)
        # amplitude
        sp_amp = torch.norm(x_stft, 2, -1).pow(2).permute(0, 2, 1).contiguous()

        # filter bank
        fb_feature = torch.log10(torch.matmul(sp_amp, self.lfcc_fb) +
                                 torch.finfo(torch.float32).eps)

        # DCT
        lfcc = self.l_dct(fb_feature)

        # Add energy
        if self.with_energy:
            power_spec = sp_amp / self.fn
            energy = torch.log10(power_spec.sum(axis=2) +
                                 torch.finfo(torch.float32).eps)
            lfcc[:, :, 0] = energy

        # Add delta coefficients
        if self.with_delta:
            lfcc_delta = delta(lfcc)
            lfcc_delta_delta = delta(lfcc_delta)
            lfcc_output = torch.cat((lfcc, lfcc_delta, lfcc_delta_delta), 2)
        else:
            lfcc_output = lfcc

        # done
        return lfcc_output

######################
### WaveForm utilities
######################

def label_2_float(x, bits):
    """Convert integer numbers to float values

    Note: dtype conversion is not handled
    Args:
    -----
       x: data to be converted Tensor.long or int, any shape.
       bits: number of bits, int

    Return:
    -------
       tensor.float

    """
    return 2 * x / (2 ** bits - 1.) - 1.


def float_2_label(x, bits):
    """Convert float wavs back to integer (quantization)

    Note: dtype conversion is not handled
    Args:
    -----
       x: data to be converted Tensor.float, any shape.
       bits: number of bits, int

    Return:
    -------
       tensor.float

    """
    # assert abs(x).max() <= 1.0
    peak = torch.abs(x).max()
    if peak > 1.0:
        x /= peak
    x = (x + 1.) * (2 ** bits - 1) / 2
    return torch.clamp(x, 0, 2 ** bits - 1)


def mulaw_encode(x, quantization_channels, scale_to_int=True):
    """Adapted from torchaudio
    https://pytorch.org/audio/functional.html mu_law_encoding
    Args:
       x (Tensor): Input tensor, float-valued waveforms in (-1, 1)
       quantization_channels (int): Number of channels
       scale_to_int: Bool
         True: scale mu-law companded to int
         False: return mu-law in (-1, 1)

    Returns:
        Tensor: Input after mu-law encoding
    """
    # mu
    mu = quantization_channels - 1.0

    # no check on the value of x
    if not x.is_floating_point():
        x = x.to(torch.float)
    mu = torch.tensor(mu, dtype=x.dtype, device=x.device)
    x_mu = torch.sign(x) * torch.log1p(mu * torch.abs(x)) / torch.log1p(mu)
    if scale_to_int:
        x_mu = ((x_mu + 1) / 2 * mu + 0.5).to(torch.int64)
    return x_mu


def mulaw_decode(x_mu, quantization_channels, input_int=True):
    """Adapted from torchaudio
    https://pytorch.org/audio/functional.html mu_law_encoding
    Args:
        x_mu (Tensor): Input tensor
        quantization_channels (int): Number of channels
    Returns:
        Tensor: Input after mu-law decoding (float-value waveform (-1, 1))
    """
    mu = quantization_channels - 1.0
    if not x_mu.is_floating_point():
        x_mu = x_mu.to(torch.float)
    mu = torch.tensor(mu, dtype=x_mu.dtype, device=x_mu.device)
    if input_int:
        x = ((x_mu) / mu) * 2 - 1.0
    else:
        x = x_mu
    x = torch.sign(x) * (torch.exp(torch.abs(x) * torch.log1p(mu)) - 1.0) / mu
    return x


######################
### DCT utilities
### https://github.com/zh217/torch-dct
### LICENSE: MIT
###
######################

def dct1(x):
    """
    Discrete Cosine Transform, Type I
    :param x: the input signal
    :return: the DCT-I of the signal over the last dimension
    """
    x_shape = x.shape
    x = x.view(-1, x_shape[-1])

    return torch.rfft(
        torch.cat([x, x.flip([1])[:, 1:-1]], dim=1), 1)[:, :, 0].view(*x_shape)


def idct1(X):
    """
    The inverse of DCT-I, which is just a scaled DCT-I
    Our definition if idct1 is such that idct1(dct1(x)) == x
    :param X: the input signal
    :return: the inverse DCT-I of the signal over the last dimension
    """
    n = X.shape[-1]
    return dct1(X) / (2 * (n - 1))


def dct(x, norm=None):
    """
    Discrete Cosine Transform, Type II (a.k.a. the DCT)
    For the meaning of the parameter `norm`, see:
    https://docs.scipy.org/doc/ scipy.fftpack.dct.html
    :param x: the input signal
    :param norm: the normalization, None or 'ortho'
    :return: the DCT-II of the signal over the last dimension
    """
    x_shape = x.shape
    N = x_shape[-1]
    x = x.contiguous().view(-1, N)

    v = torch.cat([x[:, ::2], x[:, 1::2].flip([1])], dim=1)

    Vc = torch.view_as_real(torch.fft.fft(v, dim=1))
    # print(Vc.shape)

    k = - torch.arange(N, dtype=x.dtype, device=x.device)[None, :] * np.pi / (2 * N)
    W_r = torch.cos(k)
    W_i = torch.sin(k)

    V = Vc[:, :, 0] * W_r - Vc[:, :, 1] * W_i

    if norm == 'ortho':
        V[:, 0] /= np.sqrt(N) * 2
        V[:, 1:] /= np.sqrt(N / 2) * 2

    V = 2 * V.view(*x_shape)

    return V


def idct(X, norm=None):
    """
    The inverse to DCT-II, which is a scaled Discrete Cosine Transform, Type III
    Our definition of idct is that idct(dct(x)) == x
    For the meaning of the parameter `norm`, see:
    https://docs.scipy.org/doc/ scipy.fftpack.dct.html
    :param X: the input signal
    :param norm: the normalization, None or 'ortho'
    :return: the inverse DCT-II of the signal over the last dimension
    """

    x_shape = X.shape
    N = x_shape[-1]

    X_v = X.contiguous().view(-1, x_shape[-1]) / 2

    if norm == 'ortho':
        X_v[:, 0] *= np.sqrt(N) * 2
        X_v[:, 1:] *= np.sqrt(N / 2) * 2

    k = torch.arange(x_shape[-1], dtype=X.dtype,
                     device=X.device)[None, :] * np.pi / (2 * N)
    W_r = torch.cos(k)
    W_i = torch.sin(k)

    V_t_r = X_v
    V_t_i = torch.cat([X_v[:, :1] * 0, -X_v.flip([1])[:, :-1]], dim=1)

    V_r = V_t_r * W_r - V_t_i * W_i
    V_i = V_t_r * W_i + V_t_i * W_r

    V = torch.cat([V_r.unsqueeze(2), V_i.unsqueeze(2)], dim=2)

    v = torch.irfft(V, 1, onesided=False)
    x = v.new_zeros(v.shape)
    x[:, ::2] += v[:, :N - (N // 2)]
    x[:, 1::2] += v.flip([1])[:, :N // 2]

    return x.view(*x_shape)


class LinearDCT(torch_nn.Linear):
    """Implement any DCT as a linear layer; in practice this executes around
    50x faster on GPU. Unfortunately, the DCT matrix is stored, which will
    increase memory usage.
    :param in_features: size of expected input
    :param type: which dct function in this file to use"""

    def __init__(self, in_features, type, norm=None, bias=False):
        self.type = type
        self.N = in_features
        self.norm = norm
        super(LinearDCT, self).__init__(in_features, in_features, bias=bias)

    def reset_parameters(self):
        # initialise using dct function
        I = torch.eye(self.N)
        if self.type == 'dct1':
            self.weight.data = dct1(I).data.t()
        elif self.type == 'idct1':
            self.weight.data = idct1(I).data.t()
        elif self.type == 'dct':
            self.weight.data = dct(I, norm=self.norm).data.t()
        elif self.type == 'idct':
            self.weight.data = idct(I, norm=self.norm).data.t()
        self.weight.requires_grad = False  # don't learn this!


if __name__ == "__main__":
    lfcc = LFCC(320, 160, 512, 16000, 20, with_energy=False)
    wav, sr = librosa.load("/data/neil/DS_10283_3336/LA/ASVspoof2019_LA_train/flac/LA_T_3727749.flac", sr=16000)
    # wav = torch.randn(1, 32456)
    wav = torch.Tensor(np.expand_dims(wav, axis=0))
    wav_lfcc = lfcc(wav)
    with open('/dataNVME/neil/ASVspoof2019LAFeatures/train' + '/' + "LA_T_3727749" + "LFCC" + '.pkl', 'rb') as feature_handle:
        ref_lfcc = pickle.load(feature_handle)
    print(ref_lfcc.shape)
    print(ref_lfcc[0:3,0:3])
    print(wav_lfcc.shape)
    print(wav_lfcc[0,0:3,0:3])


models/feat+resnet/audio_feature_extraction.py
yongyizang/SingFake
#!/usr/bin/python3

import numpy as np
import torch
from torch import Tensor
from torch.utils.data import Dataset, DataLoader
import pickle
import os
import librosa
from torch.utils.data.dataloader import default_collate
import torchaudio
from audio_feature_extraction import LFCC


class ASVspoof2019LA(Dataset):
    def __init__(self, path_to_audio='/data/neil/DS_10283_3336/', path_to_features='/data2/neil/ASVspoof2019LA/',
                 part='train', feature='LFCC', feat_len=400, genuine_only=False):
        super(ASVspoof2019LA, self).__init__()
        self.path_to_audio = path_to_audio
        self.path_to_features = path_to_features
        self.part = part
        self.ptf = os.path.join(path_to_features, self.part)
        self.feat_len = feat_len
        self.feature = feature
        self.genuine_only = genuine_only
        self.tag = {"-": 0, "A01": 1, "A02": 2, "A03": 3, "A04": 4, "A05": 5, "A06": 6, "A07": 7, "A08": 8, "A09": 9,
                      "A10": 10, "A11": 11, "A12": 12, "A13": 13, "A14": 14, "A15": 15, "A16": 16, "A17": 17, "A18": 18,
                      "A19": 19}
        self.label = {"spoof": 1, "bonafide": 0}
        self.all_files = librosa.util.find_files(os.path.join(self.ptf, self.feature), ext="pt")
        if self.genuine_only:
            if self.part in ["train", "dev"]:
                num_bonafide = {"train": 2580, "dev": 2548}
                self.all_files = self.all_files[:num_bonafide[self.part]]
            else:
                res = []
                for item in self.all_files:
                    if "bonafide" in item:
                        res.append(item)
                self.all_files = res
                assert len(self.all_files) == 7355

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

    def __getitem__(self, idx):
        filepath = self.all_files[idx]
        basename = os.path.basename(filepath)
        all_info = basename.split(".")[0].split("_")
        filename = "_".join(all_info[1:4])
        if self.feature != "Raw":
            featureTensor = torch.load(filepath)
            this_feat_len = featureTensor.shape[1]
            if this_feat_len > self.feat_len:
                startp = np.random.randint(this_feat_len - self.feat_len)
                featureTensor = featureTensor[:, startp:startp + self.feat_len, :]
            elif this_feat_len < self.feat_len:
                featureTensor = repeat_padding_Tensor(featureTensor, self.feat_len)
        else:
            # file_path = os.path.join(self.path_to_audio, "LA/ASVspoof2019_LA_" + self.part, "flac", filename+".flac")
            # featureTensor, sr = torchaudio.load(file_path)
            featureTensor = torch.load(filepath)
            this_feat_len = featureTensor.shape[1]
            if this_feat_len > self.feat_len:
                startp = np.random.randint(this_feat_len - self.feat_len)
                featureTensor = featureTensor[:, startp:startp + self.feat_len]
            elif this_feat_len < self.feat_len:
                featureTensor = repeat_padding_RawTensor(featureTensor, self.feat_len)
        tag = self.tag[all_info[4]]
        label = self.label[all_info[5]]
        return featureTensor, filename, tag, label, 2019

    def collate_fn(self, samples):
        return default_collate(samples)
    

class ASVspoof2019LASpec(Dataset):
    def __init__(self, path_to_audio='/data1/neil/DS_10283_3336/', path_to_features='/data2/neil/ASVspoof2019LA/',
                 part='train', feature='LFCC', feat_len=400, genuine_only=False):
        super(ASVspoof2019LASpec, self).__init__()
        self.ptd = path_to_audio
        self.part = part
        self.path_to_audio = os.path.join(self.ptd, "LA", 'ASVspoof2019_LA_'+ self.part +'/flac/')
        self.path_to_protocol = "/data1/neil/DS_10283_3336/LA/ASVspoof2019_LA_cm_protocols/"
        protocol = os.path.join(self.path_to_protocol, 'ASVspoof2019.LA.cm.'+ self.part + '.trl.txt')
        if self.part == "eval":
            protocol = os.path.join(self.ptd, "LA", 'ASVspoof2019_LA' +
                                    '_cm_protocols/ASVspoof2019.LA' + '.cm.' + self.part + '.trl.txt')
        self.tag = {"-": 0, "A01": 1, "A02": 2, "A03": 3, "A04": 4, "A05": 5, "A06": 6, "A07": 7, "A08": 8, "A09": 9,
                    "A10": 10, "A11": 11, "A12": 12, "A13": 13, "A14": 14, "A15": 15, "A16": 16, "A17": 17, "A18": 18,
                    "A19": 19}
        self.label = {"spoof": 1, "bonafide": 0}
        self.feat_len = feat_len
        self.spec = torchaudio.transforms.Spectrogram(n_fft=512, hop_length=160, win_length=512, power=2, normalized=True)

        with open(protocol, 'r') as f:
            audio_info = [info.strip().split() for info in f.readlines()]
            self.all_info = audio_info
        
    def __len__(self):
        return len(self.all_info)

    def __getitem__(self, idx):
        speaker, filename, _, tag, label = self.all_info[idx]
        filepath = os.path.join(self.path_to_audio, filename + ".flac")
        featureTensor, _ = torchaudio.load(filepath)
        this_feat_len = featureTensor.shape[1]
        if this_feat_len > self.feat_len:
            startp = np.random.randint(this_feat_len - self.feat_len)
            featureTensor = featureTensor[:, startp:startp + self.feat_len]
        elif this_feat_len < self.feat_len:
            featureTensor = repeat_padding_RawTensor(featureTensor, self.feat_len)

        featureTensor = self.spec(featureTensor)
        tag = self.tag[tag]
        label = self.label[label]
        return featureTensor, filename, tag, label, 2019

    def collate_fn(self, samples):
        return default_collate(samples)
    

class Dataset_SingFake(Dataset):
    def __init__(self, base_dir, is_mixture=False, target_sr=16000):
        """
        base_dir should contain mixtures/ and vocals/ folders
        """
        self.base_dir = base_dir
        self.is_mixture = is_mixture
        self.target_sr = target_sr
        self.cut = 64000  # take 4 sec audio (64000 samples)
        
        # get file list
        self.file_list = []
        if self.is_mixture:
            self.target_path = os.path.join(self.base_dir, "mixtures")
        else:
            self.target_path = os.path.join(self.base_dir, "vocals")
        
        assert os.path.exists(self.target_path), f"{self.target_path} does not exist!"
        
        for file in os.listdir(self.target_path):
            if file.endswith(".flac"):
                self.file_list.append(file[:-5])

        # self.lfcc = LFCC(320, 160, 512, 16000, 20, with_energy=False)
        self.spec = torchaudio.transforms.Spectrogram(n_fft=512, hop_length=160, win_length=512, power=2, normalized=True)
    
    def __len__(self):
        return len(self.file_list)

    def __getitem__(self, index):
        key = self.file_list[index]
        file_path = os.path.join(self.target_path, key + ".flac")
        # X, _ = sf.read(file_path, samplerate=self.target_sr)
        try:
            X, _ = librosa.load(file_path, sr=self.target_sr, mono=False)
            X = librosa.util.normalize(X)
        except:
            print(f"Error loading {file_path}")
            return self.__getitem__(np.random.randint(len(self.file_list)))
        if X.shape[0] > 1:
            # if not mono, take random channel
            channel_id = np.random.randint(X.shape[0])
            X = X[channel_id]
        X_pad = pad_random(X, self.cut)
        x_inp = Tensor(X_pad)
        # x_inp = self.lfcc(x_inp.unsqueeze(0))
        x_inp = self.spec(x_inp.unsqueeze(0))#.squeeze(0).transpose(0, 1)

        y = int(key.split("_")[0])
        return x_inp, y


def pad_random(x: np.ndarray, max_len: int = 64600):
    x_len = x.shape[0]
    # if duration is already long enough
    if x_len >= max_len:
        stt = np.random.randint(x_len - max_len)
        return x[stt:stt + max_len]

    # if too short
    num_repeats = int(max_len / x_len) + 1
    padded_x = np.tile(x, (num_repeats))[:max_len]
    return padded_x


class VCC2020(Dataset):
    def __init__(self, path_to_features="/data2/neil/VCC2020/", feature='LFCC',
                 feat_len=750, genuine_only=False):
        super(VCC2020, self).__init__()
        self.ptf = path_to_features
        self.feat_len = feat_len
        self.feature = feature
        self.tag = {"-": 0, "SOU": 20, "T01": 21, "T02": 22, "T03": 23, "T04": 24, "T05": 25, "T06": 26, "T07": 27, "T08": 28, "T09": 29,
                    "T10": 30, "T11": 31, "T12": 32, "T13": 33, "T14": 34, "T15": 35, "T16": 36, "T17": 37, "T18": 38, "T19": 39,
                    "T20": 40, "T21": 41, "T22": 42, "T23": 43, "T24": 44, "T25": 45, "T26": 46, "T27": 47, "T28": 48, "T29": 49,
                    "T30": 50, "T31": 51, "T32": 52, "T33": 53, "TAR": 54}
        self.label = {"spoof": 1, "bonafide": 0}
        self.genuine_only = genuine_only
        self.all_files = librosa.util.find_files(os.path.join(self.ptf, self.feature), ext="pt")

    def __len__(self):
        if self.genuine_only:
            return 220
        return len(self.all_files)

    def __getitem__(self, idx):
        filepath = self.all_files[idx]
        basename = os.path.basename(filepath)
        all_info = basename.split(".")[0].split("_")
        featureTensor = torch.load(filepath)
        this_feat_len = featureTensor.shape[1]
        if this_feat_len > self.feat_len:
            startp = np.random.randint(this_feat_len - self.feat_len)
            featureTensor = featureTensor[:, startp:startp + self.feat_len, :]
        elif this_feat_len < self.feat_len:
            featureTensor = repeat_padding_Tensor(featureTensor, self.feat_len)
        tag = self.tag[all_info[-2]]
        label = self.label[all_info[-1]]
        return featureTensor, basename, tag, label, 2020

    def collate_fn(self, samples):
        return default_collate(samples)


class ASVspoof2015(Dataset):
    def __init__(self, path_to_features, part='train', feature='LFCC', feat_len=750):
        super(ASVspoof2015, self).__init__()
        self.path_to_features = path_to_features
        self.part = part
        self.ptf = os.path.join(path_to_features, self.part)
        self.feat_len = feat_len
        self.feature = feature
        self.tag = {"human": 0, "S1": 1, "S2": 2, "S3": 3, "S4": 4, "S5": 5,
                    "S6": 6, "S7": 7, "S8": 8, "S9": 9, "S10": 10}
        self.label = {"spoof": 1, "human": 0}
        self.all_files = librosa.util.find_files(os.path.join(self.ptf, self.feature), ext="pt")

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

    def __getitem__(self, idx):
        filepath = self.all_files[idx]
        basename = os.path.basename(filepath)
        all_info = basename.split(".")[0].split("_")
        featureTensor = torch.load(filepath)
        this_feat_len = featureTensor.shape[1]
        if this_feat_len > self.feat_len:
            startp = np.random.randint(this_feat_len - self.feat_len)
            featureTensor = featureTensor[:, startp:startp + self.feat_len, :]
        elif this_feat_len < self.feat_len:
            featureTensor = repeat_padding_Tensor(featureTensor, self.feat_len)
        filename =  all_info[1]
        tag = self.tag[all_info[-2]]
        label = self.label[all_info[-1]]
        return featureTensor, filename, tag, label, 2015

    def collate_fn(self, samples):
        return default_collate(samples)


def repeat_padding_Tensor(spec, ref_len):
    mul = int(np.ceil(ref_len / spec.shape[1]))
    spec = spec.repeat(1, mul, 1)[:, :ref_len, :]
    return spec

def repeat_padding_RawTensor(raw, ref_len):
    mul = int(np.ceil(ref_len / raw.shape[1]))
    raw = raw.repeat(1, mul)[:, :ref_len]
    return raw

if __name__ == "__main__":
    dataset = Dataset_SingFake("/home/yongyi/split_0831/train", True, 16000)
    print(len(dataset))
    x, y = dataset[0]
    print(x.shape, y)



models/feat+resnet/dataset.py
yongyizang/SingFake
import numpy as np
import sys

def obtain_asv_error_rates(tar_asv, non_asv, spoof_asv, asv_threshold):

    # False alarm and miss rates for ASV
    Pfa_asv = sum(non_asv >= asv_threshold) / non_asv.size
    Pmiss_asv = sum(tar_asv < asv_threshold) / tar_asv.size

    # Rate of rejecting spoofs in ASV
    if spoof_asv.size == 0:
        Pmiss_spoof_asv = None
    else:
        Pmiss_spoof_asv = np.sum(spoof_asv < asv_threshold) / spoof_asv.size

    return Pfa_asv, Pmiss_asv, Pmiss_spoof_asv


def compute_det_curve(target_scores, nontarget_scores):

    n_scores = target_scores.size + nontarget_scores.size
    all_scores = np.concatenate((target_scores, nontarget_scores))
    labels = np.concatenate((np.ones(target_scores.size), np.zeros(nontarget_scores.size)))

    # Sort labels based on scores
    indices = np.argsort(all_scores, kind='mergesort')
    labels = labels[indices]

    # Compute false rejection and false acceptance rates
    tar_trial_sums = np.cumsum(labels)
    nontarget_trial_sums = nontarget_scores.size - (np.arange(1, n_scores + 1) - tar_trial_sums)

    frr = np.concatenate((np.atleast_1d(0), tar_trial_sums / target_scores.size))  # false rejection rates
    far = np.concatenate((np.atleast_1d(1), nontarget_trial_sums / nontarget_scores.size))  # false acceptance rates
    thresholds = np.concatenate((np.atleast_1d(all_scores[indices[0]] - 0.001), all_scores[indices]))  # Thresholds are the sorted scores

    return frr, far, thresholds


def compute_eer(target_scores, nontarget_scores):
    """ Returns equal error rate (EER) and the corresponding threshold. """
    frr, far, thresholds = compute_det_curve(target_scores, nontarget_scores)
    abs_diffs = np.abs(frr - far)
    min_index = np.argmin(abs_diffs)
    eer = np.mean((frr[min_index], far[min_index]))
    return eer, thresholds[min_index]


def compute_tDCF(bonafide_score_cm, spoof_score_cm, Pfa_asv, Pmiss_asv, Pmiss_spoof_asv, cost_model, print_cost=False):
    """
    Compute Tandem Detection Cost Function (t-DCF) [1] for a fixed ASV system.
    In brief, t-DCF returns a detection cost of a cascaded system of this form,

      Speech waveform -> [CM] -> [ASV] -> decision

    where CM stands for countermeasure and ASV for automatic speaker
    verification. The CM is therefore used as a 'gate' to decided whether or
    not the input speech sample should be passed onwards to the ASV system.
    Generally, both CM and ASV can do detection errors. Not all those errors
    are necessarily equally cost, and not all types of users are necessarily
    equally likely. The tandem t-DCF gives a principled with to compare
    different spoofing countermeasures under a detection cost function
    framework that takes that information into account.

    INPUTS:

      bonafide_score_cm   A vector of POSITIVE CLASS (bona fide or human)
                          detection scores obtained by executing a spoofing
                          countermeasure (CM) on some positive evaluation trials.
                          trial represents a bona fide case.
      spoof_score_cm      A vector of NEGATIVE CLASS (spoofing attack)
                          detection scores obtained by executing a spoofing
                          CM on some negative evaluation trials.
      Pfa_asv             False alarm (false acceptance) rate of the ASV
                          system that is evaluated in tandem with the CM.
                          Assumed to be in fractions, not percentages.
      Pmiss_asv           Miss (false rejection) rate of the ASV system that
                          is evaluated in tandem with the spoofing CM.
                          Assumed to be in fractions, not percentages.
      Pmiss_spoof_asv     Miss rate of spoof samples of the ASV system that
                          is evaluated in tandem with the spoofing CM. That
                          is, the fraction of spoof samples that were
                          rejected by the ASV system.
      cost_model          A struct that contains the parameters of t-DCF,
                          with the following fields.

                          Ptar        Prior probability of target speaker.
                          Pnon        Prior probability of nontarget speaker (zero-effort impostor)
                          Psoof       Prior probability of spoofing attack.
                          Cmiss_asv   Cost of ASV falsely rejecting target.
                          Cfa_asv     Cost of ASV falsely accepting nontarget.
                          Cmiss_cm    Cost of CM falsely rejecting target.
                          Cfa_cm      Cost of CM falsely accepting spoof.

      print_cost          Print a summary of the cost parameters and the
                          implied t-DCF cost function?

    OUTPUTS:

      tDCF_norm           Normalized t-DCF curve across the different CM
                          system operating points; see [2] for more details.
                          Normalized t-DCF > 1 indicates a useless
                          countermeasure (as the tandem system would do
                          better without it). min(tDCF_norm) will be the
                          minimum t-DCF used in ASVspoof 2019 [2].
      CM_thresholds       Vector of same size as tDCF_norm corresponding to
                          the CM threshold (operating point).

    NOTE:
    o     In relative terms, higher detection scores values are assumed to
          indicate stronger support for the bona fide hypothesis.
    o     You should provide real-valued soft scores, NOT hard decisions. The
          recommendation is that the scores are log-likelihood ratios (LLRs)
          from a bonafide-vs-spoof hypothesis based on some statistical model.
          This, however, is NOT required. The scores can have arbitrary range
          and scaling.
    o     Pfa_asv, Pmiss_asv, Pmiss_spoof_asv are in fractions, not percentages.

    References:

      [1] T. Kinnunen, K.-A. Lee, H. Delgado, N. Evans, M. Todisco,
          M. Sahidullah, J. Yamagishi, D.A. Reynolds: "t-DCF: a Detection
          Cost Function for the Tandem Assessment of Spoofing Countermeasures
          and Automatic Speaker Verification", Proc. Odyssey 2018: the
          Speaker and Language Recognition Workshop, pp. 312--319, Les Sables d'Olonne,
          France, June 2018 (https://www.isca-speech.org/archive/Odyssey_2018/pdfs/68.pdf)

      [2] ASVspoof 2019 challenge evaluation plan
          TODO: <add link>
    """


    # Sanity check of cost parameters
    if cost_model['Cfa_asv'] < 0 or cost_model['Cmiss_asv'] < 0 or \
            cost_model['Cfa_cm'] < 0 or cost_model['Cmiss_cm'] < 0:
        print('WARNING: Usually the cost values should be positive!')

    if cost_model['Ptar'] < 0 or cost_model['Pnon'] < 0 or cost_model['Pspoof'] < 0 or \
            np.abs(cost_model['Ptar'] + cost_model['Pnon'] + cost_model['Pspoof'] - 1) > 1e-10:
        sys.exit('ERROR: Your prior probabilities should be positive and sum up to one.')

    # Unless we evaluate worst-case model, we need to have some spoof tests against asv
    if Pmiss_spoof_asv is None:
        sys.exit('ERROR: you should provide miss rate of spoof tests against your ASV system.')

    # Sanity check of scores
    combined_scores = np.concatenate((bonafide_score_cm, spoof_score_cm))
    if np.isnan(combined_scores).any() or np.isinf(combined_scores).any():
        sys.exit('ERROR: Your scores contain nan or inf.')

    # Sanity check that inputs are scores and not decisions
    n_uniq = np.unique(combined_scores).size
    if n_uniq < 3:
        sys.exit('ERROR: You should provide soft CM scores - not binary decisions')

    # Obtain miss and false alarm rates of CM
    Pmiss_cm, Pfa_cm, CM_thresholds = compute_det_curve(bonafide_score_cm, spoof_score_cm)

    # Constants - see ASVspoof 2019 evaluation plan
    C1 = cost_model['Ptar'] * (cost_model['Cmiss_cm'] - cost_model['Cmiss_asv'] * Pmiss_asv) - \
         cost_model['Pnon'] * cost_model['Cfa_asv'] * Pfa_asv
    C2 = cost_model['Cfa_cm'] * cost_model['Pspoof'] * (1 - Pmiss_spoof_asv)

    # Sanity check of the weights
    if C1 < 0 or C2 < 0:
        sys.exit('You should never see this error but I cannot evalute tDCF with negative weights - please check whether your ASV error rates are correctly computed?')

    # Obtain t-DCF curve for all thresholds
    tDCF = C1 * Pmiss_cm + C2 * Pfa_cm

    # Normalized t-DCF
    tDCF_norm = tDCF / np.minimum(C1, C2)

    # Everything should be fine if reaching here.
    if print_cost:

        print('t-DCF evaluation from [Nbona={}, Nspoof={}] trials\n'.format(bonafide_score_cm.size, spoof_score_cm.size))
        print('t-DCF MODEL')
        print('   Ptar         = {:8.5f} (Prior probability of target user)'.format(cost_model['Ptar']))
        print('   Pnon         = {:8.5f} (Prior probability of nontarget user)'.format(cost_model['Pnon']))
        print('   Pspoof       = {:8.5f} (Prior probability of spoofing attack)'.format(cost_model['Pspoof']))
        print('   Cfa_asv      = {:8.5f} (Cost of ASV falsely accepting a nontarget)'.format(cost_model['Cfa_asv']))
        print('   Cmiss_asv    = {:8.5f} (Cost of ASV falsely rejecting target speaker)'.format(cost_model['Cmiss_asv']))
        print('   Cfa_cm       = {:8.5f} (Cost of CM falsely passing a spoof to ASV system)'.format(cost_model['Cfa_cm']))
        print('   Cmiss_cm     = {:8.5f} (Cost of CM falsely blocking target utterance which never reaches ASV)'.format(cost_model['Cmiss_cm']))
        print('\n   Implied normalized t-DCF function (depends on t-DCF parameters and ASV errors), s=CM threshold)')

        if C2 == np.minimum(C1, C2):
            print('   tDCF_norm(s) = {:8.5f} x Pmiss_cm(s) + Pfa_cm(s)\n'.format(C1 / C2))
        else:
            print('   tDCF_norm(s) = Pmiss_cm(s) + {:8.5f} x Pfa_cm(s)\n'.format(C2 / C1))

    return tDCF_norm, CM_thresholds

models/feat+resnet/eval_metrics.py
yongyizang/SingFake
import os
import numpy as np
import eval_metrics as em
import matplotlib.pyplot as plt

def compute_eer_and_tdcf(cm_score_file, path_to_database, visualize=False):
    asv_score_file = os.path.join(path_to_database, 'LA/ASVspoof2019_LA_asv_scores/ASVspoof2019.LA.asv.eval.gi.trl.scores.txt')

    # Fix tandem detection cost function (t-DCF) parameters
    Pspoof = 0.05
    cost_model = {
        'Pspoof': Pspoof,  # Prior probability of a spoofing attack
        'Ptar': (1 - Pspoof) * 0.99,  # Prior probability of target speaker
        'Pnon': (1 - Pspoof) * 0.01,  # Prior probability of nontarget speaker
        'Cmiss_asv': 1,  # Cost of ASV system falsely rejecting target speaker
        'Cfa_asv': 10,  # Cost of ASV system falsely accepting nontarget speaker
        'Cmiss_cm': 1,  # Cost of CM system falsely rejecting target speaker
        'Cfa_cm': 10,  # Cost of CM system falsely accepting spoof
    }

    # Load organizers' ASV scores
    asv_data = np.genfromtxt(asv_score_file, dtype=str)
    asv_sources = asv_data[:, 0]
    asv_keys = asv_data[:, 1]
    asv_scores = asv_data[:, 2].astype(np.float)

    # Load CM scores
    cm_data = np.genfromtxt(cm_score_file, dtype=str)
    cm_utt_id = cm_data[:, 0]
    cm_sources = cm_data[:, 1]
    cm_keys = cm_data[:, 2]
    cm_scores = cm_data[:, 3].astype(np.float)

    # Extract target, nontarget, and spoof scores from the ASV scores
    tar_asv = asv_scores[asv_keys == 'target']
    non_asv = asv_scores[asv_keys == 'nontarget']
    spoof_asv = asv_scores[asv_keys == 'spoof']

    # Extract bona fide (real human) and spoof scores from the CM scores
    bona_cm = cm_scores[cm_keys == 'bonafide']
    spoof_cm = cm_scores[cm_keys == 'spoof']

    # EERs of the standalone systems and fix ASV operating point to EER threshold
    eer_asv, asv_threshold = em.compute_eer(tar_asv, non_asv)
    eer_cm = em.compute_eer(bona_cm, spoof_cm)[0]

    [Pfa_asv, Pmiss_asv, Pmiss_spoof_asv] = em.obtain_asv_error_rates(tar_asv, non_asv, spoof_asv, asv_threshold)

    # Compute t-DCF
    tDCF_curve, CM_thresholds = em.compute_tDCF(bona_cm, spoof_cm, Pfa_asv, Pmiss_asv, Pmiss_spoof_asv, cost_model)

    # Minimum t-DCF
    min_tDCF_index = np.argmin(tDCF_curve)
    min_tDCF = tDCF_curve[min_tDCF_index]

    # print('ASV SYSTEM')
    # print('   EER            = {:8.5f} % (Equal error rate (target vs. nontarget discrimination)'.format(eer_asv * 100))
    # print('   Pfa            = {:8.5f} % (False acceptance rate of nontargets)'.format(Pfa_asv * 100))
    # print('   Pmiss          = {:8.5f} % (False rejection rate of targets)'.format(Pmiss_asv * 100))
    # print('   1-Pmiss,spoof  = {:8.5f} % (Spoof false acceptance rate)'.format((1 - Pmiss_spoof_asv) * 100))

    print('\nCM SYSTEM')
    print('   EER            = {:8.5f} % (Equal error rate for countermeasure)'.format(eer_cm * 100))

    print('\nTANDEM')
    print('   min-tDCF       = {:8.5f}'.format(min_tDCF))

    if visualize:
        # Visualize ASV scores and CM scores
        plt.figure()
        ax = plt.subplot(121)
        plt.hist(tar_asv, histtype='step', density=True, bins=50, label='Target')
        plt.hist(non_asv, histtype='step', density=True, bins=50, label='Nontarget')
        plt.hist(spoof_asv, histtype='step', density=True, bins=50, label='Spoof')
        plt.plot(asv_threshold, 0, 'o', markersize=10, mfc='none', mew=2, clip_on=False, label='EER threshold')
        plt.legend()
        plt.xlabel('ASV score')
        plt.ylabel('Density')
        plt.title('ASV score histogram')

        ax = plt.subplot(122)
        plt.hist(bona_cm, histtype='step', density=True, bins=50, label='Bona fide')
        plt.hist(spoof_cm, histtype='step', density=True, bins=50, label='Spoof')
        plt.legend()
        plt.xlabel('CM score')
        # plt.ylabel('Density')
        plt.title('CM score histogram')
        plt.savefig(cm_score_file[:-4]+'1.png')


        # Plot t-DCF as function of the CM threshold.
        plt.figure()
        plt.plot(CM_thresholds, tDCF_curve)
        plt.plot(CM_thresholds[min_tDCF_index], min_tDCF, 'o', markersize=10, mfc='none', mew=2)
        plt.xlabel('CM threshold index (operating point)')
        plt.ylabel('Norm t-DCF')
        plt.title('Normalized tandem t-DCF')
        plt.plot([np.min(CM_thresholds), np.max(CM_thresholds)], [1, 1], '--', color='black')
        plt.legend(('t-DCF', 'min t-DCF ({:.5f})'.format(min_tDCF), 'Arbitrarily bad CM (Norm t-DCF=1)'))
        plt.xlim([np.min(CM_thresholds), np.max(CM_thresholds)])
        plt.ylim([0, 1.5])
        plt.savefig(cm_score_file[:-4]+'2.png')

        plt.show()

    return eer_cm, min_tDCF


if __name__ == "__main__":
    # Replace CM scores with your own scores or provide score file as the first argument.
    cm_score_file = 'cm_score_ocnn.txt'
    # Replace ASV scores with organizers' scores or provide score file as the second argument.
    # path_to_database = '/home/yzh298/Downloads/DS_10283_3336/'
    path_to_database = '/data/neil/DS_10283_3336/'  # if run on GPU

    # args = sys.argv
    # if len(args) > 1:
    #     if len(args) != 3:
    #         print('USAGE: python evaluate_tDCF_asvspoof19.py <CM_SCOREFILE> <ASV_SCOREFILE>')
    #         exit()
    #     else:
    #         cm_score_file = args[1]
    #         asv_score_file = args[2]

    compute_eer_and_tdcf(cm_score_file, path_to_database)


models/feat+resnet/evaluate_tDCF_asvspoof19.py
yongyizang/SingFake
import torch
import torchaudio
import io, os
import numpy as np
from audio_feature_extraction import LFCC, repeat_padding_RawTensor
import torch.nn.functional as F


class ASVspoofHandler(object):
    def __init__(self):
        self._context = None
        self.initialized = False
        self.model = None
        self.device = None
        # Sampling rate for ASVspoof model must be 16k
        self.expected_sampling_rate = 16000
        self.feat_len = 120000

    def initialize(self, context):
        """Initialize properties and load model"""
        self._context = context
        properties = context.system_properties

        # See https://pytorch.org/serve/custom_service.html#handling-model-execution-on-multiple-gpus
        self.device = torch.device("cuda:" + str(properties.get("gpu_id")) if torch.cuda.is_available() else "cpu")

        # Read model serialize/pt file
        model_dir = properties.get("model_dir")
        serialized_file = context.manifest['model']['serializedFile']
        model_pt_path = os.path.join(model_dir, serialized_file)
        if not os.path.isfile(model_pt_path):
            raise RuntimeError("Missing the model.pt file")

        self.model = torch.jit.load(model_pt_path)

        self.initialized = True

    def _load_model(self, model_path):
        self.model = torch.load(model_path, map_location=torch.device("cpu"))
        self.model.eval()

    def preprocess(self, featureTensor):
        lfcc = LFCC(320, 160, 512, 16000, 20, with_energy=False)
        this_feat_len = featureTensor.shape[1]
        if this_feat_len > self.feat_len:
            startp = np.random.randint(this_feat_len - self.feat_len)
            featureTensor = featureTensor[:, startp:startp + self.feat_len]
        elif this_feat_len < self.feat_len:
            featureTensor = repeat_padding_RawTensor(featureTensor, self.feat_len)
        featureTensor = lfcc(featureTensor)
        return featureTensor

    def handle(self, data, context):
        """Transform input to tensor, resample, run model and return transcribed text."""
        input = data[0].get("data")
        if input is None:
            input = data[0].get("body")

        # torchaudio.load accepts file like object, here `input` is bytes
        model_input, sample_rate = torchaudio.load(io.BytesIO(input))

        # Ensure sampling rate is the same as the trained model
        if sample_rate != self.expected_sampling_rate:
            model_input = torchaudio.functional.resample(model_input, sample_rate, self.expected_sampling_rate)

        ## get input tensor is done, now doing preprocessing
        model_input = self.preprocess(model_input)

        ## inference
        feat = model_input.to(self.device).unsqueeze(0).transpose(2, 3)
        feats, feat_outputs = self.model(feat)

        # score0 = F.softmax(feat_outputs)[:, 0].item()
        # score1 = F.softmax(feat_outputs)[:, 1].item()

        score0 = feat_outputs[:, 0].item()

        # ## get output
        # if score0 > score1:
        #     return ["bonafide\n"]
        # else:
        #     return ["spoofing\n"]
        return ["{\"score\": %.3f}\n" %score0]

models/feat+resnet/handler.py
yongyizang/SingFake
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import math


class OCSoftmax(nn.Module):
    def __init__(self, feat_dim=2, m_real=0.5, m_fake=0.2, alpha=20.0):
        super(OCSoftmax, self).__init__()
        self.feat_dim = feat_dim
        self.m_real = m_real
        self.m_fake = m_fake
        self.alpha = alpha
        self.center = nn.Parameter(torch.randn(1, self.feat_dim))
        nn.init.kaiming_uniform_(self.center, 0.25)
        self.softplus = nn.Softplus()

    def forward(self, x, labels):
        """
        Args:
            x: feature matrix with shape (batch_size, feat_dim).
            labels: ground truth labels with shape (batch_size).
        """
        w = F.normalize(self.center, p=2, dim=1)
        x = F.normalize(x, p=2, dim=1)

        scores = x @ w.transpose(0,1)
        output_scores = scores.clone()

        scores[labels == 0] = self.m_real - scores[labels == 0]
        scores[labels == 1] = scores[labels == 1] - self.m_fake

        loss = self.softplus(self.alpha * scores).mean()

        return loss, output_scores.squeeze(1)


class AMSoftmax(nn.Module):
    def __init__(self, num_classes, enc_dim, s=20, m=0.5):
        super(AMSoftmax, self).__init__()
        self.enc_dim = enc_dim
        self.num_classes = num_classes
        self.s = s
        self.m = m
        self.centers = nn.Parameter(torch.randn(num_classes, enc_dim))

    def forward(self, feat, label):
        batch_size = feat.shape[0]
        norms = torch.norm(feat, p=2, dim=-1, keepdim=True)
        nfeat = torch.div(feat, norms)

        norms_c = torch.norm(self.centers, p=2, dim=-1, keepdim=True)
        ncenters = torch.div(self.centers, norms_c)
        logits = torch.matmul(nfeat, torch.transpose(ncenters, 0, 1))

        y_onehot = torch.FloatTensor(batch_size, self.num_classes)
        y_onehot.zero_()
        y_onehot = Variable(y_onehot).cuda()
        y_onehot.scatter_(1, torch.unsqueeze(label, dim=-1), self.m)
        margin_logits = self.s * (logits - y_onehot)

        return logits, margin_logits


class IsolateLoss(nn.Module):
    """Isolate loss.
        Reference:
        I. Masi, A. Killekar, R. M. Mascarenhas, S. P. Gurudatt, and W. AbdAlmageed, “Two-branch Recurrent Network for Isolating Deepfakes in Videos,” 2020, [Online]. Available: http://arxiv.org/abs/2008.03412.
        Args:
            num_classes (int): number of classes.
            feat_dim (int): feature dimension.
            r_real (float): small radius to keep real inside
            r_fake (float): large radius to keep fake outside
        """
    def __init__(self, num_classes=2, feat_dim=256, r_real=25, r_fake=75):
        super(IsolateLoss, self).__init__()
        self.num_classes = num_classes
        self.feat_dim = feat_dim
        self.r_real = r_real
        self.r_fake = r_fake

        self.center = nn.Parameter(torch.randn(1, self.feat_dim))

    def forward(self, x, labels):
        """
        Args:
            x: feature matrix with shape (batch_size, feat_dim).
            labels: ground truth labels with shape (batch_size).
        """
        loss = F.relu(torch.norm(x[labels==0]-self.center, p=2, dim=1) - self.r_real).mean() \
               + F.relu(self.r_fake - torch.norm(x[labels==1]-self.center, p=2, dim=1)).mean()
        return loss, -torch.norm(x-self.center, p=2, dim=1)


class SingleCenterLoss(nn.Module):
    """Single-Center loss.
        Reference:
        Li, J., Xie, H., Li, J., Wang, Z., & Zhang, Y. (2021). Frequency-aware Discriminative Feature Learning Supervised by Single-Center Loss for Face Forgery Detection. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 6458-6467).
        Args:
            num_classes (int): number of classes.
            feat_dim (int): feature dimension.
            m (float): scale factor that the margin is proportional to the square root of dimension
        """
    def __init__(self, num_classes=2, feat_dim=256, m=0.3):
        super(SingleCenterLoss, self).__init__()
        self.num_classes = num_classes
        self.feat_dim = feat_dim
        self.m = m

        self.center = nn.Parameter(torch.randn(1, self.feat_dim))

    def forward(self, x, labels):
        """
        Args:
            x: feature matrix with shape (batch_size, feat_dim).
            labels: ground truth labels with shape (batch_size).
        """
        m_nat = torch.norm(x[labels==0]-self.center, p=2, dim=1).mean()
        m_man = torch.norm(x[labels==1]-self.center, p=2, dim=1).mean()
        loss = m_nat + F.relu(m_nat - m_man + self.m * math.sqrt(self.feat_dim))
        return loss, -torch.norm(x-self.center, p=2, dim=1)


class AngularIsoLoss(nn.Module):
    def __init__(self, feat_dim=2, m_real=0.5, m_fake=0.2, alpha=20.0):
        super(AngularIsoLoss, self).__init__()
        self.feat_dim = feat_dim
        self.m_real = m_real
        self.m_fake = m_fake
        self.alpha = alpha
        self.center = nn.Parameter(torch.randn(1, self.feat_dim))
        nn.init.kaiming_uniform_(self.center, 0.25)
        self.softplus = nn.Softplus()

    def forward(self, x, labels):
        """
        Args:
            x: feature matrix with shape (batch_size, feat_dim).
            labels: ground truth labels with shape (batch_size).
        """
        w = F.normalize(self.center, p=2, dim=1)
        x = F.normalize(x, p=2, dim=1)

        scores = x @ w.transpose(0,1)
        output_scores = scores.clone()

        scores[labels == 0] = self.m_real - scores[labels == 0]
        scores[labels == 1] = scores[labels == 1] - self.m_fake

        loss = self.softplus(self.alpha * scores[labels == 0]).mean() + \
               self.softplus(self.alpha * scores[labels == 1]).mean()

        return loss, output_scores.squeeze(1)

models/feat+resnet/loss.py
yongyizang/SingFake
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.init as init
from torch.autograd import Function
from torch import Tensor
import os
import numpy as np


## Adapted from https://github.com/joaomonteirof/e2e_antispoofing
## https://github.com/nii-yamagishilab/project-NN-Pytorch-scripts/blob/newfunctions/


class SelfAttention(nn.Module):
    def __init__(self, hidden_size, mean_only=False):
        super(SelfAttention, self).__init__()

        #self.output_size = output_size
        self.hidden_size = hidden_size
        self.att_weights = nn.Parameter(torch.Tensor(1, hidden_size),requires_grad=True)

        self.mean_only = mean_only

        init.kaiming_uniform_(self.att_weights)

    def forward(self, inputs):
        batch_size = inputs.size(0)
        weights = torch.bmm(inputs, self.att_weights.permute(1, 0).unsqueeze(0).repeat(batch_size, 1, 1))

        if inputs.size(0)==1:
            attentions = F.softmax(torch.tanh(weights),dim=1)
            weighted = torch.mul(inputs, attentions.expand_as(inputs))
        else:
            attentions = F.softmax(torch.tanh(weights.squeeze()),dim=1)
            weighted = torch.mul(inputs, attentions.unsqueeze(2).expand_as(inputs))

        if self.mean_only:
            return weighted.sum(1)
        else:
            noise = 1e-5*torch.randn(weighted.size())

            if inputs.is_cuda:
                noise = noise.to(inputs.device)
            avg_repr, std_repr = weighted.sum(1), (weighted+noise).std(1)

            representations = torch.cat((avg_repr,std_repr),1)

            return representations


class PreActBlock(nn.Module):
    '''Pre-activation version of the BasicBlock.'''
    expansion = 1

    def __init__(self, in_planes, planes, stride, *args, **kwargs):
        super(PreActBlock, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_planes)
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)

        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = nn.Sequential(nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False))

    def forward(self, x):
        out = F.relu(self.bn1(x))
        shortcut = self.shortcut(out) if hasattr(self, 'shortcut') else x
        out = self.conv1(out)
        out = self.conv2(F.relu(self.bn2(out)))
        out += shortcut
        return out


class PreActBottleneck(nn.Module):
    '''Pre-activation version of the original Bottleneck module.'''
    expansion = 4

    def __init__(self, in_planes, planes, stride, *args, **kwargs):
        super(PreActBottleneck, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_planes)
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, self.expansion*planes, kernel_size=1, bias=False)

        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = nn.Sequential(nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False))

    def forward(self, x):
        out = F.relu(self.bn1(x))
        shortcut = self.shortcut(out) if hasattr(self, 'shortcut') else x
        out = self.conv1(out)
        out = self.conv2(F.relu(self.bn2(out)))
        out = self.conv3(F.relu(self.bn3(out)))
        out += shortcut
        return out

def conv3x3(in_planes, out_planes, stride=1):
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)

def conv1x1(in_planes, out_planes, stride=1):
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)

RESNET_CONFIGS = {'18': [[2, 2, 2, 2], PreActBlock],
                  '28': [[3, 4, 6, 3], PreActBlock],
                  '34': [[3, 4, 6, 3], PreActBlock],
                  '50': [[3, 4, 6, 3], PreActBottleneck],
                  '101': [[3, 4, 23, 3], PreActBottleneck]
                  }

class ResNet(nn.Module):
    def __init__(self, num_nodes, enc_dim, resnet_type='18', nclasses=2):
        self.in_planes = 16
        super(ResNet, self).__init__()

        layers, block = RESNET_CONFIGS[resnet_type]

        self._norm_layer = nn.BatchNorm2d

        self.conv1 = nn.Conv2d(1, 16, kernel_size=(9, 3), stride=(3, 1), padding=(1, 1), bias=False)
        self.bn1 = nn.BatchNorm2d(16)
        self.activation = nn.ReLU()

        self.layer1 = self._make_layer(block, 64, layers[0], stride=1)
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)

        self.conv5 = nn.Conv2d(512 * block.expansion, 256, kernel_size=(num_nodes, 3), stride=(1, 1), padding=(0, 1),
                               bias=False)
        self.bn5 = nn.BatchNorm2d(256)

        self.conv6 = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=(7, 3), stride=(1, 1))
        self.bn6 = nn.BatchNorm2d(256)
        

        self.fc = nn.Linear(256 * 2, enc_dim)
        self.fc_mu = nn.Linear(enc_dim, nclasses) if nclasses >= 2 else nn.Linear(enc_dim, 1)

        self.initialize_params()
        self.attention = SelfAttention(256)

    def initialize_params(self):
        for layer in self.modules():
            if isinstance(layer, torch.nn.Conv2d):
                init.kaiming_normal_(layer.weight, a=0, mode='fan_out')
            elif isinstance(layer, torch.nn.Linear):
                init.kaiming_uniform_(layer.weight)
            elif isinstance(layer, torch.nn.BatchNorm2d) or isinstance(layer, torch.nn.BatchNorm1d):
                layer.weight.data.fill_(1)
                layer.bias.data.zero_()

    def _make_layer(self, block, planes, num_blocks, stride=1):
        norm_layer = self._norm_layer
        downsample = None
        if stride != 1 or self.in_planes != planes * block.expansion:
            downsample = nn.Sequential(conv1x1(self.in_planes, planes * block.expansion, stride),
                                       norm_layer(planes * block.expansion))
        layers = []
        layers.append(block(self.in_planes, planes, stride, downsample, 1, 64, 1, norm_layer))
        self.in_planes = planes * block.expansion
        for _ in range(1, num_blocks):
            layers.append(
                block(self.in_planes, planes, 1, groups=1, base_width=64, dilation=False, norm_layer=norm_layer))

        return nn.Sequential(*layers)

    def forward(self, x):

        x = self.conv1(x)
        # print(x.shape)
        x = self.activation(self.bn1(x))
        # print(x.shape)
        x = self.layer1(x)
        # print(x.shape)
        x = self.layer2(x)
        # print(x.shape)
        x = self.layer3(x)
        # print(x.shape)
        x = self.layer4(x)
        # print(x.shape)
        x = self.conv5(x)
        # print(x.shape)
        
        # x = self.activation(self.bn5(x)).squeeze(2)

        x = self.activation(self.bn5(x))
        x = self.conv6(x)
        x = self.activation(self.bn6(x)).squeeze(2)

        # print(x.shape)

        stats = self.attention(x.permute(0, 2, 1).contiguous())

        feat = self.fc(stats)

        mu = self.fc_mu(feat)

        return feat, mu


class MaxFeatureMap2D(nn.Module):
    """ Max feature map (along 2D)

    MaxFeatureMap2D(max_dim=1)

    l_conv2d = MaxFeatureMap2D(1)
    data_in = torch.rand([1, 4, 5, 5])
    data_out = l_conv2d(data_in)

    Input:
    ------
    data_in: tensor of shape (batch, channel, ...)

    Output:
    -------
    data_out: tensor of shape (batch, channel//2, ...)

    Note
    ----
    By default, Max-feature-map is on channel dimension,
    and maxout is used on (channel ...)
    """

    def __init__(self, max_dim=1):
        super(MaxFeatureMap2D, self).__init__()
        self.max_dim = max_dim

    def forward(self, inputs):
        # suppose inputs (batchsize, channel, length, dim)

        shape = list(inputs.size())

        if self.max_dim >= len(shape):
            print("MaxFeatureMap: maximize on %d dim" % (self.max_dim))
            print("But input has %d dimensions" % (len(shape)))
            sys.exit(1)
        if shape[self.max_dim] // 2 * 2 != shape[self.max_dim]:
            print("MaxFeatureMap: maximize on %d dim" % (self.max_dim))
            print("But this dimension has an odd number of data")
            sys.exit(1)
        shape[self.max_dim] = shape[self.max_dim] // 2
        shape.insert(self.max_dim, 2)

        # view to (batchsize, 2, channel//2, ...)
        # maximize on the 2nd dim
        m, i = inputs.view(*shape).max(self.max_dim)
        return m


class LCNN(nn.Module):
    def __init__(self, num_nodes, args, nclasses=2):
        super(LCNN, self).__init__()
        self.num_nodes = num_nodes
        self.enc_dim = args.enc_dim
        self.nclasses = nclasses
        self.conv1 = nn.Sequential(nn.Conv2d(1, 64, (5, 5), 1, padding=(2, 2)),
                                   MaxFeatureMap2D(),
                                   nn.MaxPool2d((2, 2), (2, 2)))
        self.conv2 = nn.Sequential(nn.Conv2d(32, 64, (1, 1), 1, padding=(0, 0)),
                                   MaxFeatureMap2D(),
                                   nn.BatchNorm2d(32, affine=False))
        self.conv3 = nn.Sequential(nn.Conv2d(32, 96, (3, 3), 1, padding=(1, 1)),
                                   MaxFeatureMap2D(),
                                   nn.MaxPool2d((2, 2), (2, 2)),
                                   nn.BatchNorm2d(48, affine=False))
        self.conv4 = nn.Sequential(nn.Conv2d(48, 96, (1, 1), 1, padding=(0, 0)),
                                   MaxFeatureMap2D(),
                                   nn.BatchNorm2d(48, affine=False))
        self.conv5 = nn.Sequential(nn.Conv2d(48, 128, (3, 3), 1, padding=(1, 1)),
                                   MaxFeatureMap2D(),
                                   nn.MaxPool2d((2, 2), (2, 2)))
        self.conv6 = nn.Sequential(nn.Conv2d(64, 128, (1, 1), 1, padding=(0, 0)),
                                   MaxFeatureMap2D(),
                                   nn.BatchNorm2d(64, affine=False))
        self.conv7 = nn.Sequential(nn.Conv2d(64, 64, (3, 3), 1, padding=(1, 1)),
                                   MaxFeatureMap2D(),
                                   nn.BatchNorm2d(32, affine=False))
        self.conv8 = nn.Sequential(nn.Conv2d(32, 64, (1, 1), 1, padding=(0, 0)),
                                   MaxFeatureMap2D(),
                                   nn.BatchNorm2d(32, affine=False))
        self.conv9 = nn.Sequential(nn.Conv2d(32, 64, (3, 3), 1, padding=[1, 1]),
                                   MaxFeatureMap2D(),
                                   nn.MaxPool2d((2, 2), (2, 2)))
        self.out = nn.Sequential(nn.Dropout(0.7),
                                 nn.Linear((args.feat_len // 16) * (60 // 16) * 32, 160),
                                 MaxFeatureMap2D(),
                                 nn.Linear(80, self.enc_dim))
        self.fc_mu = nn.Linear(self.enc_dim, nclasses) if nclasses >= 2 else nn.Linear(self.enc_dim, 1)

    def forward(self, x):

        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        x = self.conv6(x)
        x = self.conv7(x)
        x = self.conv8(x)
        x = self.conv9(x)
        feat = torch.flatten(x, 1)
        feat = self.out(feat)
        out = self.fc_mu(feat)

        return feat, out

class GradientReversalFunction(Function):
    """
    Gradient Reversal Layer from:
    Unsupervised Domain Adaptation by Backpropagation (Ganin & Lempitsky, 2015)
    Forward pass is the identity function. In the backward pass,
    the upstream gradients are multiplied by -lambda (i.e. gradient is reversed)
    """

    @staticmethod
    def forward(ctx, x, lambda_):
        ctx.lambda_ = lambda_
        return x.clone()

    @staticmethod
    def backward(ctx, grads):
        lambda_ = ctx.lambda_
        lambda_ = grads.new_tensor(lambda_)
        dx = -lambda_ * grads
        return dx, None


class GradientReversal(nn.Module):
    def __init__(self, lambda_=1):
        super(GradientReversal, self).__init__()
        self.lambda_ = lambda_

    def forward(self, x):
        return GradientReversalFunction.apply(x, self.lambda_)


class ChannelClassifier(nn.Module):
    def __init__(self, enc_dim, nclasses, lambda_=0.05, ADV=True):
        super(ChannelClassifier, self).__init__()
        self.adv = ADV
        if self.adv:
            self.grl = GradientReversal(lambda_)
        self.classifier = nn.Sequential(nn.Linear(enc_dim, enc_dim // 2),
                                        nn.Dropout(0.3),
                                        nn.ReLU(),
                                        nn.Linear(enc_dim // 2, nclasses),
                                        nn.ReLU())

    def initialize_params(self):
        for layer in self.modules():
            if isinstance(layer, torch.nn.Linear):
                init.kaiming_uniform_(layer.weight)

    def forward(self, x):
        if self.adv:
            x = self.grl(x)
        return self.classifier(x)


class SincConv(nn.Module):
    @staticmethod
    def to_mel(hz):
        return 2595 * np.log10(1 + hz / 700)

    @staticmethod
    def to_hz(mel):
        return 700 * (10 ** (mel / 2595) - 1)

    def __init__(self, device, out_channels, kernel_size, in_channels=1, sample_rate=16000,
                 stride=1, padding=0, dilation=1, bias=False, groups=1):

        super(SincConv, self).__init__()

        if in_channels != 1:
            msg = "SincConv only support one input channel (here, in_channels = {%i})" % (in_channels)
            raise ValueError(msg)

        self.device = device
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.sample_rate = sample_rate

        # Forcing the filters to be odd (i.e, perfectly symmetrics)
        if kernel_size % 2 == 0:
            self.kernel_size = self.kernel_size + 1

        self.stride = stride
        self.padding = padding
        self.dilation = dilation

        if bias:
            raise ValueError('SincConv does not support bias.')
        if groups > 1:
            raise ValueError('SincConv does not support groups.')

        # initialize filterbanks using Mel scale
        NFFT = 512
        f = int(self.sample_rate / 2) * np.linspace(0, 1, int(NFFT / 2) + 1)
        fmel = self.to_mel(f)  # Hz to mel conversion
        fmelmax = np.max(fmel)
        fmelmin = np.min(fmel)
        filbandwidthsmel = np.linspace(fmelmin, fmelmax, self.out_channels + 1)
        filbandwidthsf = self.to_hz(filbandwidthsmel)  # Mel to Hz conversion
        self.mel = filbandwidthsf
        self.hsupp = torch.arange(-(self.kernel_size - 1) / 2, (self.kernel_size - 1) / 2 + 1)
        self.band_pass = torch.zeros(self.out_channels, self.kernel_size)

    def forward(self, x):
        for i in range(len(self.mel) - 1):
            fmin = self.mel[i]
            fmax = self.mel[i + 1]
            hHigh = (2 * fmax / self.sample_rate) * np.sinc(2 * fmax * self.hsupp / self.sample_rate)
            hLow = (2 * fmin / self.sample_rate) * np.sinc(2 * fmin * self.hsupp / self.sample_rate)
            hideal = hHigh - hLow

            self.band_pass[i, :] = Tensor(np.hamming(self.kernel_size)) * Tensor(hideal)

        band_pass_filter = self.band_pass.to(self.device)

        self.filters = (band_pass_filter).view(self.out_channels, 1, self.kernel_size)

        return F.conv1d(x, self.filters, stride=self.stride,
                        padding=self.padding, dilation=self.dilation,
                        bias=None, groups=1)


class Residual_block(nn.Module):
    def __init__(self, nb_filts, first=False):
        super(Residual_block, self).__init__()
        self.first = first

        if not self.first:
            self.bn1 = nn.BatchNorm1d(num_features=nb_filts[0])

        self.lrelu = nn.LeakyReLU(negative_slope=0.3)

        self.conv1 = nn.Conv1d(in_channels=nb_filts[0],
                               out_channels=nb_filts[1],
                               kernel_size=3,
                               padding=1,
                               stride=1)

        self.bn2 = nn.BatchNorm1d(num_features=nb_filts[1])
        self.conv2 = nn.Conv1d(in_channels=nb_filts[1],
                               out_channels=nb_filts[1],
                               padding=1,
                               kernel_size=3,
                               stride=1)

        if nb_filts[0] != nb_filts[1]:
            self.downsample = True
            self.conv_downsample = nn.Conv1d(in_channels=nb_filts[0],
                                             out_channels=nb_filts[1],
                                             padding=0,
                                             kernel_size=1,
                                             stride=1)

        else:
            self.downsample = False
        self.mp = nn.MaxPool1d(3)

    def forward(self, x):
        identity = x
        if not self.first:
            out = self.bn1(x)
            out = self.lrelu(out)
        else:
            out = x

        out = self.conv1(x)
        out = self.bn2(out)
        out = self.lrelu(out)
        out = self.conv2(out)

        if self.downsample:
            identity = self.conv_downsample(identity)

        out += identity
        out = self.mp(out)
        return out


class RawNet(nn.Module):
    def __init__(self, d_args, args):
        super(RawNet, self).__init__()

        self.device = args.device

        self.Sinc_conv = SincConv(device=self.device,
                                  out_channels=d_args['filts'][0],
                                  kernel_size=d_args['first_conv'],
                                  in_channels=d_args['in_channels']
                                  )

        self.first_bn = nn.BatchNorm1d(num_features=d_args['filts'][0])
        self.selu = nn.SELU(inplace=True)
        self.block0 = nn.Sequential(Residual_block(nb_filts=d_args['filts'][1], first=True))
        self.block1 = nn.Sequential(Residual_block(nb_filts=d_args['filts'][1]))
        self.block2 = nn.Sequential(Residual_block(nb_filts=d_args['filts'][2]))
        d_args['filts'][2][0] = d_args['filts'][2][1]
        self.block3 = nn.Sequential(Residual_block(nb_filts=d_args['filts'][2]))
        self.block4 = nn.Sequential(Residual_block(nb_filts=d_args['filts'][2]))
        self.block5 = nn.Sequential(Residual_block(nb_filts=d_args['filts'][2]))
        self.avgpool = nn.AdaptiveAvgPool1d(1)

        self.fc_attention0 = self._make_attention_fc(in_features=d_args['filts'][1][-1],
                                                     l_out_features=d_args['filts'][1][-1])
        self.fc_attention1 = self._make_attention_fc(in_features=d_args['filts'][1][-1],
                                                     l_out_features=d_args['filts'][1][-1])
        self.fc_attention2 = self._make_attention_fc(in_features=d_args['filts'][2][-1],
                                                     l_out_features=d_args['filts'][2][-1])
        self.fc_attention3 = self._make_attention_fc(in_features=d_args['filts'][2][-1],
                                                     l_out_features=d_args['filts'][2][-1])
        self.fc_attention4 = self._make_attention_fc(in_features=d_args['filts'][2][-1],
                                                     l_out_features=d_args['filts'][2][-1])
        self.fc_attention5 = self._make_attention_fc(in_features=d_args['filts'][2][-1],
                                                     l_out_features=d_args['filts'][2][-1])

        self.bn_before_gru = nn.BatchNorm1d(num_features=d_args['filts'][2][-1])
        self.gru = nn.GRU(input_size=d_args['filts'][2][-1],
                          hidden_size=d_args['gru_node'],
                          num_layers=d_args['nb_gru_layer'],
                          batch_first=True)

        self.fc1_gru = nn.Linear(in_features=d_args['gru_node'],
                                 out_features=args.enc_dim)

        self.fc2_gru = nn.Linear(in_features=args.enc_dim,
                                 out_features=d_args['nb_classes'], bias=True)

        self.sig = nn.Sigmoid()
        self.logsoftmax = nn.LogSoftmax(dim=1)

    def forward(self, x, y=None):

        x = self.Sinc_conv(x)
        x = F.max_pool1d(torch.abs(x), 3)
        x = self.first_bn(x)
        x = self.selu(x)

        x0 = self.block0(x)
        y0 = self.avgpool(x0).view(x0.size(0), -1)  # torch.Size([batch, filter])
        y0 = self.fc_attention0(y0)
        y0 = self.sig(y0).view(y0.size(0), y0.size(1), -1)  # torch.Size([batch, filter, 1])
        x = x0 * y0 + y0  # (batch, filter, time) x (batch, filter, 1)

        x1 = self.block1(x)
        y1 = self.avgpool(x1).view(x1.size(0), -1)  # torch.Size([batch, filter])
        y1 = self.fc_attention1(y1)
        y1 = self.sig(y1).view(y1.size(0), y1.size(1), -1)  # torch.Size([batch, filter, 1])
        x = x1 * y1 + y1  # (batch, filter, time) x (batch, filter, 1)

        x2 = self.block2(x)
        y2 = self.avgpool(x2).view(x2.size(0), -1)  # torch.Size([batch, filter])
        y2 = self.fc_attention2(y2)
        y2 = self.sig(y2).view(y2.size(0), y2.size(1), -1)  # torch.Size([batch, filter, 1])
        x = x2 * y2 + y2  # (batch, filter, time) x (batch, filter, 1)

        x3 = self.block3(x)
        y3 = self.avgpool(x3).view(x3.size(0), -1)  # torch.Size([batch, filter])
        y3 = self.fc_attention3(y3)
        y3 = self.sig(y3).view(y3.size(0), y3.size(1), -1)  # torch.Size([batch, filter, 1])
        x = x3 * y3 + y3  # (batch, filter, time) x (batch, filter, 1)

        x4 = self.block4(x)
        y4 = self.avgpool(x4).view(x4.size(0), -1)  # torch.Size([batch, filter])
        y4 = self.fc_attention4(y4)
        y4 = self.sig(y4).view(y4.size(0), y4.size(1), -1)  # torch.Size([batch, filter, 1])
        x = x4 * y4 + y4  # (batch, filter, time) x (batch, filter, 1)

        x5 = self.block5(x)
        y5 = self.avgpool(x5).view(x5.size(0), -1)  # torch.Size([batch, filter])
        y5 = self.fc_attention5(y5)
        y5 = self.sig(y5).view(y5.size(0), y5.size(1), -1)  # torch.Size([batch, filter, 1])
        x = x5 * y5 + y5  # (batch, filter, time) x (batch, filter, 1)

        x = self.bn_before_gru(x)
        x = self.selu(x)
        x = x.permute(0, 2, 1)  # (batch, filt, time) >> (batch, time, filt)
        self.gru.flatten_parameters()
        x, _ = self.gru(x)
        x = x[:, -1, :]
        feat = self.fc1_gru(x)
        output = self.fc2_gru(feat)
        # output = self.logsoftmax(x)

        return feat, output

    def _make_attention_fc(self, in_features, l_out_features):

        l_fc = []

        l_fc.append(nn.Linear(in_features=in_features,
                              out_features=l_out_features))

        return nn.Sequential(*l_fc)

    def _make_layer(self, nb_blocks, nb_filts, first=False):
        layers = []
        # def __init__(self, nb_filts, first = False):
        for i in range(nb_blocks):
            first = first if i == 0 else False
            layers.append(Residual_block(nb_filts=nb_filts,
                                         first=first))
            if i == 0: nb_filts[0] = nb_filts[1]

        return nn.Sequential(*layers)

    def summary(self, input_size, batch_size=-1, device="cuda", print_fn=None):
        if print_fn == None: printfn = print
        model = self

        def register_hook(module):
            def hook(module, input, output):
                class_name = str(module.__class__).split(".")[-1].split("'")[0]
                module_idx = len(summary)

                m_key = "%s-%i" % (class_name, module_idx + 1)
                summary[m_key] = OrderedDict()
                summary[m_key]["input_shape"] = list(input[0].size())
                summary[m_key]["input_shape"][0] = batch_size
                if isinstance(output, (list, tuple)):
                    summary[m_key]["output_shape"] = [
                        [-1] + list(o.size())[1:] for o in output
                    ]
                else:
                    summary[m_key]["output_shape"] = list(output.size())
                    if len(summary[m_key]["output_shape"]) != 0:
                        summary[m_key]["output_shape"][0] = batch_size

                params = 0
                if hasattr(module, "weight") and hasattr(module.weight, "size"):
                    params += torch.prod(torch.LongTensor(list(module.weight.size())))
                    summary[m_key]["trainable"] = module.weight.requires_grad
                if hasattr(module, "bias") and hasattr(module.bias, "size"):
                    params += torch.prod(torch.LongTensor(list(module.bias.size())))
                summary[m_key]["nb_params"] = params

            if (
                    not isinstance(module, nn.Sequential)
                    and not isinstance(module, nn.ModuleList)
                    and not (module == model)
            ):
                hooks.append(module.register_forward_hook(hook))

        device = device.lower()
        assert device in [
            "cuda",
            "cpu",
        ], "Input device is not valid, please specify 'cuda' or 'cpu'"

        if device == "cuda" and torch.cuda.is_available():
            dtype = torch.cuda.FloatTensor
        else:
            dtype = torch.FloatTensor
        if isinstance(input_size, tuple):
            input_size = [input_size]
        x = [torch.rand(2, *in_size).type(dtype) for in_size in input_size]
        summary = OrderedDict()
        hooks = []
        model.apply(register_hook)
        model(*x)
        for h in hooks:
            h.remove()

        print_fn("----------------------------------------------------------------")
        line_new = "{:>20}  {:>25} {:>15}".format("Layer (type)", "Output Shape", "Param #")
        print_fn(line_new)
        print_fn("================================================================")
        total_params = 0
        total_output = 0
        trainable_params = 0
        for layer in summary:
            # input_shape, output_shape, trainable, nb_params
            line_new = "{:>20}  {:>25} {:>15}".format(
                layer,
                str(summary[layer]["output_shape"]),
                "{0:,}".format(summary[layer]["nb_params"]),
            )
            total_params += summary[layer]["nb_params"]
            total_output += np.prod(summary[layer]["output_shape"])
            if "trainable" in summary[layer]:
                if summary[layer]["trainable"] == True:
                    trainable_params += summary[layer]["nb_params"]
            print_fn(line_new)


if __name__ == "__main__":
    os.environ["CUDA_VISIBLE_DEVICES"] = "1"

    # cqcc = torch.randn((32,1,90,788)).cuda()
    # resnet = ResNet(4, 2, resnet_type='18', nclasses=2).cuda()
    # _, output = resnet(cqcc)
    # print(output.shape)
    lfcc = torch.randn((1, 1, 60, 750)).cuda()
    lcnn = LCNN(4, 2, nclasses=2).cuda()
    feat, output = lcnn(lfcc)
    print(output.shape)
    # cnn = ConvNet(num_classes = 2, num_nodes = 47232, enc_dim = 256).cuda()
    # _, output = cnn(lfcc)
    # print(output.shape)

models/feat+resnet/model.py
yongyizang/SingFake
import raw_dataset as dataset
from audio_feature_extraction import LFCC
import os
import torch
from tqdm import tqdm
from torchaudio import transforms
from scipy.fftpack import fft, ifft, fftshift, ifftshift, next_fast_len
import numpy as np
import copy

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

cuda = torch.cuda.is_available()
print('Cuda device available: ', cuda)
device = torch.device("cuda" if cuda else "cpu")


# ## LFCC
# for part_ in ["train", "dev", "eval"]:
#     asvspoof_raw = dataset.ASVspoof2019Raw("LA", "/data1/neil/DS_10283_3336/", "/data1/neil/DS_10283_3336/LA/ASVspoof2019_LA_cm_protocols/", part=part_)
#     target_dir = os.path.join("/data2/neil/ASVspoof2019LA", part_, "LFCC")
#     if not os.path.exists(target_dir):
#         os.makedirs(target_dir)
#     lfcc = LFCC(320, 160, 512, 16000, 20, with_energy=False)
#     lfcc = lfcc.to(device)
#     for idx in tqdm(range(len(asvspoof_raw))):
#         # if idx > 0: break # debug purpose
#         waveform, filename, tag, label = asvspoof_raw[idx]
#         # waveform = spectral_whitening(waveform.squeeze(0).numpy())
#         # waveform = torch.from_numpy(np.expand_dims(waveform, axis=0))
#         waveform = waveform.to(device)
#         lfccOfWav = lfcc(waveform)
#         # print(lfccOfWav.shape)
#         torch.save(lfccOfWav, os.path.join(target_dir, "%05d_%s_%s_%s.pt" % (idx, filename, tag, label)))
#     print("Done!")

## MFCC
for part_ in ["train", "dev", "eval"]:
    asvspoof_raw = dataset.ASVspoof2019Raw("LA", "/data1/neil/DS_10283_3336/", "/data1/neil/DS_10283_3336/LA/ASVspoof2019_LA_cm_protocols/", part=part_)
    target_dir = os.path.join("/data2/neil/ASVspoof2019LA", part_, "MFCC")
    if not os.path.exists(target_dir):
        os.makedirs(target_dir)
    transform = transforms.MFCC(sample_rate=16000, n_mfcc=60, melkwargs={"n_fft": 320, "hop_length": 160, "center": False})
    transform = transform.to(device)
    for idx in tqdm(range(len(asvspoof_raw))):
        # if idx > 0: break
        waveform, filename, tag, label = asvspoof_raw[idx]
        # waveform = spectral_whitening(waveform.squeeze(0).numpy())
        # waveform = torch.from_numpy(np.expand_dims(waveform, axis=0))
        waveform = waveform.to(device)
        mfccOfWav = transform(waveform).transpose(2, 1)
        # print(mfccOfWav.shape)
        torch.save(mfccOfWav, os.path.join(target_dir, "%05d_%s_%s_%s.pt" % (idx, filename, tag, label)))
    print("Done!")

models/feat+resnet/preprocess.py
yongyizang/SingFake
#!/usr/bin/python3

import numpy as np
import torch
from torch import Tensor
from torch.utils.data import Dataset
import scipy.io as sio
import pickle
import os
import librosa
from torch.utils.data.dataloader import default_collate
import warnings


def torchaudio_load(filepath):
    wave, sr = librosa.load(filepath, sr=16000)
    wave = librosa.util.normalize(wave)
    waveform = torch.Tensor(np.expand_dims(wave, axis=0))
    return [waveform, sr]

class ASVspoof2019Raw(Dataset):
    def __init__(self, access_type, path_to_database, path_to_protocol, part='train'):
        super(ASVspoof2019Raw, self).__init__()
        self.access_type = access_type
        self.ptd = path_to_database
        self.part = part
        self.path_to_audio = os.path.join(self.ptd, access_type, 'ASVspoof2019_'+access_type+'_'+ self.part +'/flac/')
        self.path_to_protocol = path_to_protocol
        protocol = os.path.join(self.path_to_protocol, 'ASVspoof2019.'+access_type+'.cm.'+ self.part + '.trl.txt')
        if self.part == "eval":
            protocol = os.path.join(self.ptd, access_type, 'ASVspoof2019_' + access_type +
                                    '_cm_protocols/ASVspoof2019.' + access_type + '.cm.' + self.part + '.trl.txt')
        if self.access_type == 'LA':
            self.tag = {"-": 0, "A01": 1, "A02": 2, "A03": 3, "A04": 4, "A05": 5, "A06": 6, "A07": 7, "A08": 8, "A09": 9,
                      "A10": 10, "A11": 11, "A12": 12, "A13": 13, "A14": 14, "A15": 15, "A16": 16, "A17": 17, "A18": 18,
                      "A19": 19}
        else:
            self.tag = {"-": 0, "AA": 1, "AB": 2, "AC": 3, "BA": 4, "BB": 5, "BC": 6, "CA": 7, "CB": 8, "CC": 9}
        self.label = {"spoof": 1, "bonafide": 0}

        with open(protocol, 'r') as f:
            audio_info = [info.strip().split() for info in f.readlines()]
            self.all_info = audio_info

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

    def __getitem__(self, idx):
        speaker, filename, _, tag, label = self.all_info[idx]
        filepath = os.path.join(self.path_to_audio, filename + ".flac")
        waveform, sr = torchaudio_load(filepath)

        return waveform, filename, tag, label

    def collate_fn(self, samples):
        return default_collate(samples)


class VCC2020Raw(Dataset):
    def __init__(self, path_to_spoof="/data2/neil/nii-yamagishilab-VCC2020-listeningtest-31f913c", path_to_bonafide="/data2/neil/nii-yamagishilab-VCC2020-database-0b2fb2e"):
        super(VCC2020Raw, self).__init__()
        self.all_spoof = librosa.util.find_files(path_to_spoof, ext="wav")
        self.all_bonafide = librosa.util.find_files(path_to_bonafide, ext="wav")

    def __len__(self):
        # print(len(self.all_spoof), len(self.all_bonafide))
        return len(self.all_spoof) + len(self.all_bonafide)

    def __getitem__(self, idx):
        if idx < len(self.all_bonafide):
            filepath = self.all_bonafide[idx]
            label = "bonafide"
            filename = "_".join(filepath.split("/")[-3:])[:-4]
            tag = "-"
        else:
            filepath = self.all_spoof[idx - len(self.all_bonafide)]
            filename = os.path.basename(filepath)[:-4]
            label = "spoof"
            tag = filepath.split("/")[-3]
        waveform, sr = torchaudio_load(filepath)

        return waveform, filename, tag, label

    def collate_fn(self, samples):
        return default_collate(samples)


class ASVspoof2015Raw(Dataset):
    def __init__(self, path_to_database="/data/neil/ASVspoof2015/wav", path_to_protocol="/data/neil/ASVspoof2015/CM_protocol", part='train'):
        super(ASVspoof2015Raw, self).__init__()
        self.ptd = path_to_database
        self.part = part
        self.path_to_audio = os.path.join(self.ptd, self.part)
        self.path_to_protocol = path_to_protocol
        cm_pro_dict = {"train": "cm_train.trn", "dev": "cm_develop.ndx", "eval": "cm_evaluation.ndx"}
        protocol = os.path.join(self.path_to_protocol, cm_pro_dict[self.part])
        self.tag = {"human": 0, "S1": 1, "S2": 2, "S3": 3, "S4": 4, "S5": 5,
                    "S6": 6, "S7": 7, "S8": 8, "S9": 9, "S10": 10}
        self.label = {"spoof": 1, "human": 0}

        with open(protocol, 'r') as f:
            audio_info = [info.strip().split() for info in f.readlines()]
            self.all_info = audio_info

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

    def __getitem__(self, idx):
        speaker, filename, tag, label = self.all_info[idx]
        filepath = os.path.join(self.path_to_audio, speaker, filename + ".wav")
        waveform, sr = torchaudio_load(filepath)
        filename = filename.replace("_", "-")
        return waveform, filename, tag, label

    def collate_fn(self, samples):
        return default_collate(samples)


class ASVspoof2019LARaw_withChannel(Dataset):
    def __init__(self, access_type="LA", path_to_database="/data/shared/ASVspoof2019Channel", path_to_protocol="/data/neil/DS_10283_3336/LA/ASVspoof2019_LA_cm_protocols/", part='train'):
        super(ASVspoof2019LARaw_withChannel, self).__init__()
        self.access_type = access_type
        self.ptd = path_to_database
        self.part = part
        self.path_to_audio = path_to_database
        self.path_to_protocol = path_to_protocol
        protocol = os.path.join(self.path_to_protocol,
                                'ASVspoof2019.' + access_type + '.cm.' + self.part + '.trl.txt')
        if self.part == "eval":
            protocol = os.path.join(self.ptd, access_type, 'ASVspoof2019_' + access_type +
                                    '_cm_protocols/ASVspoof2019.' + access_type + '.cm.' + self.part + '.trl.txt')
        self.tag = {"-": 0, "A01": 1, "A02": 2, "A03": 3, "A04": 4, "A05": 5, "A06": 6, "A07": 7, "A08": 8,
                    "A09": 9,
                    "A10": 10, "A11": 11, "A12": 12, "A13": 13, "A14": 14, "A15": 15, "A16": 16, "A17": 17,
                    "A18": 18,
                    "A19": 19}
        self.label = {"spoof": 1, "bonafide": 0}
        self.channel = ['amr[br=5k15]', 'amrwb[br=15k85]', 'g711[law=u]', 'g722[br=56k]',
                        'g722[br=64k]', 'g726[law=a,br=16k]', 'g728', 'g729a', 'gsmfr',
                        'silk[br=20k]', 'silk[br=5k]', 'silkwb[br=10k,loss=5]', 'silkwb[br=30k]']

        with open(protocol, 'r') as f:
            audio_info = [info.strip().split() for info in f.readlines()]
            self.all_info = audio_info

    def __len__(self):
        return len(self.all_info) * len(self.channel)

    def __getitem__(self, idx):
        file_idx = idx // len(self.channel)
        channel_idx = idx % len(self.channel)
        speaker, filename, _, tag, label = self.all_info[file_idx]
        channel = self.channel[channel_idx]
        filepath = os.path.join(self.path_to_audio, filename + "_" + channel + ".wav")
        waveform, sr = torchaudio_load(filepath)

        return waveform, filename, tag, label, channel

    def collate_fn(self, samples):
        return default_collate(samples)


class ASVspoof2019LARaw_withDevice(Dataset):
    def __init__(self, access_type="LA", path_to_database="/data/shared/ASVspoof2019LA-Sim", path_to_protocol="/data/neil/DS_10283_3336/LA/ASVspoof2019_LA_cm_protocols/", part='eval'):
        super(ASVspoof2019LARaw_withDevice, self).__init__()
        self.access_type = access_type
        self.ptd = path_to_database
        self.part = part
        self.path_to_audio = path_to_database
        self.path_to_protocol = path_to_protocol
        protocol = os.path.join(self.path_to_protocol,
                                'ASVspoof2019.' + access_type + '.cm.' + self.part + '.trl.txt')
        self.tag = {"-": 0, "A01": 1, "A02": 2, "A03": 3, "A04": 4, "A05": 5, "A06": 6, "A07": 7, "A08": 8,
                    "A09": 9,
                    "A10": 10, "A11": 11, "A12": 12, "A13": 13, "A14": 14, "A15": 15, "A16": 16, "A17": 17,
                    "A18": 18,
                    "A19": 19}
        self.label = {"spoof": 1, "bonafide": 0}
        self.devices = ['AKSPKRS80sUk002-16000', 'AKSPKRSVinUk002-16000', 'Doremi-16000', 'RCAPB90-16000',
                        'ResloRBRedLabel-16000', 'AKSPKRSSpeaker002-16000', 'BehritoneirRecording-16000',
                        'OktavaML19-16000', 'ResloRB250-16000', 'SonyC37Fet-16000', 'iPadirRecording-16000', 'iPhoneirRecording-16000']

        with open(protocol, 'r') as f:
            audio_info = [info.strip().split() for info in f.readlines()]
            self.all_info = audio_info

    def __len__(self):
        return len(self.all_info) * len(self.devices)

    def __getitem__(self, idx):
        file_idx = idx // len(self.devices)
        device_idx = idx % len(self.devices)
        speaker, filename, _, tag, label = self.all_info[file_idx]
        device = self.devices[device_idx]
        filepath = os.path.join(self.path_to_audio, device, filename + ".wav")
        waveform, sr = torchaudio_load(filepath)

        return waveform, filename, tag, label, device

    def collate_fn(self, samples):
        return default_collate(samples)

if __name__ == "__main__":

    asvspoof2019channel = ASVspoof2019LARaw_withChannel()
    print(len(asvspoof2019channel))
    waveform, filename, tag, label, channel = asvspoof2019channel[123]
    print(waveform.shape)
    print(filename)
    print(tag)
    print(label)
    print(channel)
    pass

models/feat+resnet/raw_dataset.py
yongyizang/SingFake
import logging
import os
import torch
from torch.utils.model_zoo import tqdm
import random
import numpy as np
from dataset import *
from torch.utils.data import DataLoader
import torch.nn.functional as F
import eval_metrics as em
from evaluate_tDCF_asvspoof19 import compute_eer_and_tdcf
from utils import setup_seed
import argparse
from utils import str2bool

## Adapted from https://github.com/pytorch/audio/tree/master/torchaudio
## https://github.com/nii-yamagishilab/project-NN-Pytorch-scripts/blob/newfunctions/

def init():
    parser = argparse.ArgumentParser("load model scores")
    parser.add_argument('--seed', type=int, help="random number seed", default=1000)
    parser.add_argument("-d", "--path_to_database", type=str, help="dataset path",
                        default='/data1/neil/DS_10283_3336/')
    parser.add_argument("-f", "--path_to_features", type=str, help="features path",
                        default='/data2/neil/ASVspoof2019LA/')
    parser.add_argument('-m', '--model_dir', type=str, help="directory for pretrained model", required=True,
                        default='/data3/neil/chan/adv1010')
    parser.add_argument("-t", "--task", type=str, help="which dataset you would like to test on",
                        required=True, default='ASVspoof2019LA',
                        choices=["ASVspoof2019LA", "ASVspoof2015", "VCC2020", "ASVspoof2019LASim", "ASVspoof2021LA", "singfake"])
    parser.add_argument('-l', '--loss', type=str, default="ocsoftmax",
                        choices=["softmax", "amsoftmax", "ocsoftmax", "isolate", "scl", "angulariso"],
                        help="loss for scoring")
    parser.add_argument('--weight_loss', type=float, default=0.5, help="weight for other loss")
    parser.add_argument("--feat", type=str, help="which feature to use", default='LFCC',
                        choices=["CQCC", "LFCC", "Raw", "MFCC"])
    parser.add_argument("--feat_len", type=int, help="features length", default=500)
    parser.add_argument('--batch_size', type=int, default=64, help="Mini batch size for training")
    parser.add_argument('--is_mixture', type=str2bool, nargs='?', const=True, default=False,
                        help="whether use mixture or vocals in training")
    parser.add_argument("--gpu", type=str, help="GPU index", default="0")
    args = parser.parse_args()

    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
    setup_seed(args.seed)      # Set seeds
    args.cuda = torch.cuda.is_available()
    args.device = torch.device("cuda" if args.cuda else "cpu")

    return args

def test_model_on_ASVspoof2019LA(feat_model_path, loss_model_path, part, add_loss):
    dirname = os.path.dirname
    basename = os.path.splitext(os.path.basename(feat_model_path))[0]
    if "checkpoint" in dirname(feat_model_path):
        dir_path = dirname(dirname(feat_model_path))
    else:
        dir_path = dirname(feat_model_path)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = torch.load(feat_model_path)
    loss_model = torch.load(loss_model_path) if add_loss is not None else None
    test_set = ASVspoof2019LASpec(args.path_to_database, args.path_to_features, part,
                            args.feat, feat_len=args.feat_len)
    testDataLoader = DataLoader(test_set, batch_size=args.batch_size, shuffle=False, num_workers=0)
    model.eval()
    score_loader, idx_loader = [], []

    with open(os.path.join(dir_path, 'checkpoint_cm_score_ASVspoof2019LA.txt'), 'w') as cm_score_file:
        for i, (feat, audio_fn, tags, labels, _) in enumerate(tqdm(testDataLoader)):
            if args.feat == "Raw":
                feat = feat.to(args.device)
            else:
                feat = feat.transpose(2, 3).to(args.device)
            # print(feat.shape)
            tags = tags.to(device)
            labels = labels.to(device)

            feats, feat_outputs = model(feat)

            if add_loss == "softmax":
                score = F.softmax(feat_outputs)[:, 0]
            elif add_loss == "ocsoftmax":
                ang_isoloss, score = loss_model(feats, labels)
            elif add_loss == "isolate":
                _, score = loss_model(feats, labels)
            elif add_loss == "scl":
                score_softmax = F.softmax(feat_outputs)[:, 0]
                _, score_scl = loss_model(feats, labels)
                score = score_softmax + args.weight_loss * score_scl
            elif add_loss == "amsoftmax":
                outputs, moutputs = loss_model(feats, labels)
                score = F.softmax(outputs, dim=1)[:, 0]
            elif add_loss == "angulariso":
                angularisoloss, score = loss_model(feats, labels)
            else:
                raise ValueError("what is the loss?")

            for j in range(labels.size(0)):
                cm_score_file.write(
                    '%s A%02d %s %s\n' % (audio_fn[j], tags[j].data,
                                          "spoof" if labels[j].data.cpu().numpy() else "bonafide",
                                          score[j].item()))

            # score_loader.append(score.detach().cpu())
            # idx_loader.append(labels.detach().cpu())

    # scores = torch.cat(score_loader, 0).data.cpu().numpy()
    # labels = torch.cat(idx_loader, 0).data.cpu().numpy()
    # eer = em.compute_eer(scores[labels == 0], scores[labels == 1])[0]

    eer, min_tDCF = compute_eer_and_tdcf(os.path.join(dir_path, 'checkpoint_cm_score_ASVspoof2019LA.txt'),
                                            "/data1/neil/DS_10283_3336/")

    return eer, min_tDCF

def test_model_on_singfake(feat_model_path, loss_model_path, dataset_path, add_loss, args):
    model = torch.load(feat_model_path)
    loss_model = torch.load(loss_model_path) if add_loss is not None else None
    test_set = Dataset_SingFake(dataset_path, is_mixture=args.is_mixture)
    testDataLoader = DataLoader(test_set, batch_size=args.batch_size, shuffle=False, num_workers=0)
    model.eval()
    with torch.no_grad():
        idx_loader, score_loader = [], []
        for i, (feat, labels) in enumerate(tqdm(testDataLoader)):
            if args.feat == "Raw":
                feat = feat.to(args.device)
            else:
                feat = feat.transpose(2, 3).to(args.device)

            labels = labels.to(args.device)
            feats, feat_outputs = model(feat)

            if add_loss == "softmax":
                score = F.softmax(feat_outputs)[:, 0]
            elif add_loss == "ocsoftmax":
                ang_isoloss, score = loss_model(feats, labels)
            elif add_loss == "isolate":
                _, score = loss_model(feats, labels)
            elif add_loss == "scl":
                score_softmax = F.softmax(feat_outputs)[:, 0]
                _, score_scl = loss_model(feats, labels)
                score = score_softmax + args.weight_loss * score_scl
            elif add_loss == "amsoftmax":
                outputs, moutputs = loss_model(feats, labels)
                score = F.softmax(outputs, dim=1)[:, 0]
            elif add_loss == "angulariso":
                angularisoloss, score = loss_model(feats, labels)
            else:
                raise ValueError("what is the loss?")
            

            idx_loader.append((labels))
            score_loader.append(score)

        scores = torch.cat(score_loader, 0).data.cpu().numpy()
        labels = torch.cat(idx_loader, 0).data.cpu().numpy()
        eer = em.compute_eer(scores[labels == 0], scores[labels == 1])[0]
        # print eer in percentage
        print(eer * 100)

    return eer


def test_on_VCC(feat_model_path, loss_model_path, part, add_loss):
    dirname = os.path.dirname
    basename = os.path.splitext(os.path.basename(feat_model_path))[0]
    if "checkpoint" in dirname(feat_model_path):
        dir_path = dirname(dirname(feat_model_path))
    else:
        dir_path = dirname(feat_model_path)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = torch.load(feat_model_path)
    loss_model = torch.load(loss_model_path) if add_loss is not None else None
    test_set_VCC = VCC2020("/data2/neil/VCC2020/", "LFCC", feat_len=args.feat_len)
    testDataLoader = DataLoader(test_set_VCC, batch_size=args.batch_size, shuffle=False, num_workers=0)
    model.eval()
    score_loader, idx_loader = [], []

    with open(os.path.join(dir_path, 'checkpoint_cm_score_VCC.txt'), 'w') as cm_score_file:
        for i, (feat, _, tags, labels, _) in enumerate(tqdm(testDataLoader)):
            if args.feat == "Raw":
                feat = feat.to(args.device)
            else:
                feat = feat.transpose(2, 3).to(args.device)

            tags = tags.to(device)
            labels = labels.to(device)

            feats, feat_outputs = model(feat)

            if add_loss == "softmax":
                score = F.softmax(feat_outputs)[:, 0]
            elif add_loss == "ocsoftmax":
                ang_isoloss, score = loss_model(feats, labels)
            elif add_loss == "isolate":
                _, score = loss_model(feats, labels)
            elif add_loss == "scl":
                score_softmax = F.softmax(feat_outputs)[:, 0]
                _, score_scl = loss_model(feats, labels)
                score = score_softmax + args.weight_loss * score_scl
            elif add_loss == "amsoftmax":
                outputs, moutputs = loss_model(feats, labels)
                score = F.softmax(outputs, dim=1)[:, 0]
            elif add_loss == "angulariso":
                angularisoloss, score = loss_model(feats, labels)
            else:
                raise ValueError("what is the loss?")

            for j in range(labels.size(0)):
                cm_score_file.write(
                    'A%02d %s %s\n' % (tags[j].data,
                                          "spoof" if labels[j].data.cpu().numpy() else "bonafide",
                                          score[j].item()))

            score_loader.append(score.detach().cpu())
            idx_loader.append(labels.detach().cpu())

    scores = torch.cat(score_loader, 0).data.cpu().numpy()
    labels = torch.cat(idx_loader, 0).data.cpu().numpy()
    eer = em.compute_eer(scores[labels == 0], scores[labels == 1])[0]

    return eer

def test_on_ASVspoof2015(feat_model_path, loss_model_path, part, add_loss):
    dirname = os.path.dirname
    basename = os.path.splitext(os.path.basename(feat_model_path))[0]
    if "checkpoint" in dirname(feat_model_path):
        dir_path = dirname(dirname(feat_model_path))
    else:
        dir_path = dirname(feat_model_path)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = torch.load(feat_model_path)
    loss_model = torch.load(loss_model_path) if add_loss is not None else None
    test_set_2015 = ASVspoof2015("/data2/neil/ASVspoof2015/", part="eval", feature="LFCC", feat_len=args.feat_len)
    testDataLoader = DataLoader(test_set_2015, batch_size=args.batch_size, shuffle=False, num_workers=0)
    model.eval()
    score_loader, idx_loader = [], []

    with open(os.path.join(dir_path, 'checkpoint_cm_score_ASVspoof2015.txt'), 'w') as cm_score_file:
        for i, (feat, audio_fn, tags, labels, _) in enumerate(tqdm(testDataLoader)):
            if args.feat == "Raw":
                feat = feat.to(args.device)
            else:
                feat = feat.transpose(2, 3).to(args.device)
            tags = tags.to(device)
            labels = labels.to(device)

            feats, feat_outputs = model(feat)

            if add_loss == "softmax":
                score = F.softmax(feat_outputs)[:, 0]
            elif add_loss == "ocsoftmax":
                ang_isoloss, score = loss_model(feats, labels)
            elif add_loss == "isolate":
                _, score = loss_model(feats, labels)
            elif add_loss == "scl":
                score_softmax = F.softmax(feat_outputs)[:, 0]
                _, score_scl = loss_model(feats, labels)
                score = score_softmax + args.weight_loss * score_scl
            elif add_loss == "amsoftmax":
                outputs, moutputs = loss_model(feats, labels)
                score = F.softmax(outputs, dim=1)[:, 0]
            elif add_loss == "angulariso":
                angularisoloss, score = loss_model(feats, labels)
            else:
                raise ValueError("what is the loss?")

            for j in range(labels.size(0)):
                cm_score_file.write(
                    '%s A%02d %s %s\n' % (audio_fn[j], tags[j].data,
                                          "spoof" if labels[j].data.cpu().numpy() else "bonafide",
                                          score[j].item()))

            score_loader.append(score.detach().cpu())
            idx_loader.append(labels.detach().cpu())

    scores = torch.cat(score_loader, 0).data.cpu().numpy()
    labels = torch.cat(idx_loader, 0).data.cpu().numpy()
    eer = em.compute_eer(scores[labels == 0], scores[labels == 1])[0]

    return eer

def test_individual_attacks(cm_score_file):
    # Load CM scores
    cm_data = np.genfromtxt(cm_score_file, dtype=str)
    cm_sources = cm_data[:, 1]
    cm_keys = cm_data[:, 2]
    cm_scores = cm_data[:, 3].astype(np.float)

    eer_cm_lst, min_tDCF_lst = [], []
    for attack_idx in range(0, 55):
        # Extract target, nontarget, and spoof scores from the ASV scores

        # Extract bona fide (real human) and spoof scores from the CM scores
        bona_cm = cm_scores[cm_keys == 'bonafide']
        spoof_cm = cm_scores[cm_sources == 'A%02d' % attack_idx]

        # EERs of the standalone systems and fix ASV operating point to EER threshold
        eer_cm = em.compute_eer(bona_cm, spoof_cm)[0]
        eer_cm_lst.append(eer_cm)

    return eer_cm_lst

def test_on_ASVspoof2019LASim(feat_model_path, loss_model_path, part, add_loss):
    dirname = os.path.dirname
    # basename = os.path.splitext(os.path.basename(feat_model_path))[0]
    if "checkpoint" in dirname(feat_model_path):
        dir_path = dirname(dirname(feat_model_path))
    else:
        dir_path = dirname(feat_model_path)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = torch.load(feat_model_path)
    loss_model = torch.load(loss_model_path) if add_loss is not None else None
    test_set = ASVspoof2019LASim(path_to_features="/data2/neil/ASVspoof2019LA/",
                                                path_to_deviced="/dataNVME/neil/ASVspoof2019LADevice",
                                                part="eval",
                                                feature=args.feat, feat_len=args.feat_len)
    testDataLoader = DataLoader(test_set, batch_size=args.batch_size, shuffle=True, num_workers=0)
    model.eval()
    # score_loader, idx_loader = [], []

    with open(os.path.join(dir_path, 'checkpoint_cm_score_ASVspoof2019LASim.txt'), 'w') as cm_score_file:
        for i, (feat, audio_fn, tags, labels, _) in enumerate(tqdm(testDataLoader)):
            if i > int(len(test_set) / args.batch_size / (len(test_set.devices) + 1)): break
            if args.feat == "Raw":
                feat = feat.to(args.device)
            else:
                feat = feat.transpose(2, 3).to(args.device)
            # print(feat.shape)
            tags = tags.to(device)
            labels = labels.to(device)

            feats, feat_outputs = model(feat)

            if add_loss == "softmax":
                score = F.softmax(feat_outputs)[:, 0]
            elif add_loss == "ocsoftmax":
                ang_isoloss, score = loss_model(feats, labels)
            elif add_loss == "isolate":
                _, score = loss_model(feats, labels)
            elif add_loss == "scl":
                score_softmax = F.softmax(feat_outputs)[:, 0]
                _, score_scl = loss_model(feats, labels)
                score = score_softmax + args.weight_loss * score_scl
            elif add_loss == "amsoftmax":
                outputs, moutputs = loss_model(feats, labels)
                score = F.softmax(outputs, dim=1)[:, 0]
            elif add_loss == "angulariso":
                angularisoloss, score = loss_model(feats, labels)
            else:
                raise ValueError("what is the loss?")

            for j in range(labels.size(0)):
                cm_score_file.write(
                    '%s A%02d %s %s\n' % (audio_fn[j], tags[j].data,
                                          "spoof" if labels[j].data.cpu().numpy() else "bonafide",
                                          score[j].item()))

    #         score_loader.append(score.detach().cpu())
    #         idx_loader.append(labels.detach().cpu())
    #
    # scores = torch.cat(score_loader, 0).data.cpu().numpy()
    # labels = torch.cat(idx_loader, 0).data.cpu().numpy()
    # eer = em.compute_eer(scores[labels == 0], scores[labels == 1])[0]

    eer, min_tDCF = compute_eer_and_tdcf(os.path.join(dir_path, 'checkpoint_cm_score_ASVspoof2019LASim.txt'),
                                            "/data/neil/DS_10283_3336/")

    return eer, min_tDCF


def test_on_ASVspoof2021LA(feat_model_path, loss_model_path, part, add_loss):
    dirname = os.path.dirname
    if "checkpoint" in dirname(feat_model_path):
        dir_path = dirname(dirname(feat_model_path))
    else:
        dir_path = dirname(feat_model_path)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = torch.load(feat_model_path)
    loss_model = torch.load(loss_model_path) if add_loss is not None else None

    ### use this line to generate score for LA 2021 Challenge
    test_set = ASVspoof2021LAeval(feature=args.feat, feat_len=args.feat_len)
    testDataLoader = DataLoader(test_set, batch_size=args.batch_size, shuffle=False, num_workers=0)
    model.eval()

    txt_file_name = os.path.join(dir_path, 'score.txt')

    with open(txt_file_name, 'w') as cm_score_file:
        for i, data_slice in enumerate(tqdm(testDataLoader)):
            feat, audio_fn = data_slice
            if args.feat == "Raw":
                feat = feat.to(args.device)
            else:
                feat = feat.transpose(2, 3).to(args.device)

            labels = torch.zeros((feat.shape[0]))
            labels = labels.to(device)

            feats, feat_outputs = model(feat)

            if add_loss == "softmax":
                score = F.softmax(feat_outputs)[:, 0]
            elif add_loss == "ocsoftmax":
                ang_isoloss, score = loss_model(feats, labels)
            elif add_loss == "isolate":
                _, score = loss_model(feats, labels)
            elif add_loss == "scl":
                score_softmax = F.softmax(feat_outputs)[:, 0]
                _, score_scl = loss_model(feats, labels)
                score = score_softmax + args.weight_loss * score_scl
            elif add_loss == "amsoftmax":
                outputs, moutputs = loss_model(feats, labels)
                score = F.softmax(outputs, dim=1)[:, 0]
            elif add_loss == "angulariso":
                angularisoloss, score = loss_model(feats, labels)
            else:
                raise ValueError("what is the loss?")

            for j in range(labels.size(0)):
                cm_score_file.write('%s %s\n' % (audio_fn[j], score[j].item()))


if __name__ == "__main__":
    device = torch.device("cuda")

    args = init()
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu

    model_path = os.path.join(args.model_dir, "anti-spoofing_feat_model.pt")
    loss_model_path = os.path.join(args.model_dir, "anti-spoofing_loss_model.pt")

    if args.task == "ASVspoof2019LA":
        eer = test_model_on_ASVspoof2019LA(model_path, loss_model_path, "eval", args.loss)
    elif args.task == "ASVspoof2015":
        eer = test_on_ASVspoof2015(model_path, loss_model_path, "eval", args.loss)
        print(eer)
    elif args.task =="VCC2020":
        eer = test_on_VCC(model_path, loss_model_path, "eval", args.loss)
        print(eer)
    elif args.task =="ASVspoof2019LASim":
        eer = test_on_ASVspoof2019LASim(model_path, loss_model_path, "eval", args.loss)
    elif args.task == "ASVspoof2021LA":
        eer = test_on_ASVspoof2021LA(model_path, loss_model_path, "eval", args.loss)
    elif args.task == "singfake":
        sets = ["train", "dev", "test", "additional_test", "codec_test/adts_64k", "codec_test/mp3_128k", "codec_test/ogg_64k", "codec_test/opus_64k", "persian_test"]
        for set in sets:
            print(set)
            eer = test_model_on_singfake(model_path, loss_model_path, os.path.join(args.path_to_database, set), args.loss, args)
    else:
        raise ValueError("Evaluation task unknown!")



models/feat+resnet/test.py
yongyizang/SingFake
import torch
import argparse
import os
import json
import shutil
import numpy as np
from model import *
from dataset import *
from torch.utils.data import DataLoader
from loss import *
from collections import defaultdict
from tqdm import tqdm
from utils import str2bool, setup_seed
import eval_metrics as em
import yaml

torch.set_default_tensor_type(torch.FloatTensor)

def initParams():
    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument('--seed', type=int, help="random number seed", default=1000)

    # Data folder prepare
    parser.add_argument("-d", "--path_to_database", type=str, help="dataset path",
                        default='/data1/neil/DS_10283_3336/')
    parser.add_argument("-f", "--path_to_features", type=str, help="features path",
                        default='/data2/neil/ASVspoof2019LA/')
    parser.add_argument("-p", "--path_to_protocol", type=str, help="protocol path",
                        default='/data/neil/DS_10283_3336/LA/ASVspoof2019_LA_cm_protocols/')
    parser.add_argument("-o", "--out_fold", type=str, help="output folder", required=True, default='./models/try/')

    # Dataset prepare
    parser.add_argument("--feat", type=str, help="which feature to use", default='LFCC',
                        choices=["CQCC", "LFCC", "Raw", "MFCC"])
    parser.add_argument("--feat_len", type=int, help="features length", default=64000)
    parser.add_argument("--enc_dim", type=int, help="encoding dimension", default=256)

    parser.add_argument('-m', '--model', help='Model arch', default='resnet',
                        choices=['resnet', 'lcnn', 'rawnet'])

    # Training hyperparameters
    parser.add_argument('--num_epochs', type=int, default=100, help="Number of epochs for training")
    parser.add_argument('--batch_size', type=int, default=128, help="Mini batch size for training")
    parser.add_argument('--lr', type=float, default=0.0003, help="learning rate")
    parser.add_argument('--lr_decay', type=float, default=0.5, help="decay learning rate")
    parser.add_argument('--interval', type=int, default=20, help="interval to decay lr")

    parser.add_argument('--beta_1', type=float, default=0.9, help="bata_1 for Adam")
    parser.add_argument('--beta_2', type=float, default=0.999, help="beta_2 for Adam")
    parser.add_argument('--eps', type=float, default=1e-8, help="epsilon for Adam")
    parser.add_argument("--gpu", type=str, help="GPU index", default="1")
    parser.add_argument('--num_workers', type=int, default=0, help="number of workers")

    parser.add_argument('-l', '--loss', type=str, default="ocsoftmax",
                        choices=["softmax", "amsoftmax", "ocsoftmax", "isolate", "scl", "angulariso"], help="loss for training")
    parser.add_argument('--weight_loss', type=float, default=0.5, help="weight for other loss")
    parser.add_argument('--m_real', type=float, default=0.9, help="m_real for ocsoftmax loss")
    parser.add_argument('--m_fake', type=float, default=0.2, help="m_fake for ocsoftmax loss")
    parser.add_argument('--r_real', type=float, default=25.0, help="r_real for isolate loss")
    parser.add_argument('--r_fake', type=float, default=75.0, help="r_fake for isolate loss")
    parser.add_argument('--alpha', type=float, default=20, help="scale factor for amsoftmax and ocsoftmax loss")
    parser.add_argument('--scale_factor', type=float, default=0.5, help="scale factor for single center loss")

    parser.add_argument('--continue_training', action='store_true', help="continue training with trained model")

    parser.add_argument('--test_on_eval', action='store_true',
                        help="whether to run EER on the evaluation set")
    parser.add_argument('--test_interval', type=int, default=5, help="test on eval for every how many epochs")
    parser.add_argument('--save_interval', type=int, default=5, help="save checkpoint model for every how many epochs")

    args = parser.parse_args()

    # Change this to specify GPU
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu

    # Set seeds
    setup_seed(args.seed)

    if args.continue_training:
        pass
    else:
        # Path for output data
        if not os.path.exists(args.out_fold):
            os.makedirs(args.out_fold)
        else:
            shutil.rmtree(args.out_fold)
            os.mkdir(args.out_fold)

        # Folder for intermediate results
        if not os.path.exists(os.path.join(args.out_fold, 'checkpoint')):
            os.makedirs(os.path.join(args.out_fold, 'checkpoint'))
        else:
            shutil.rmtree(os.path.join(args.out_fold, 'checkpoint'))
            os.mkdir(os.path.join(args.out_fold, 'checkpoint'))

        # Path for input data
        # assert os.path.exists(args.path_to_database)
        assert os.path.exists(args.path_to_features)

        # Save training arguments
        with open(os.path.join(args.out_fold, 'args.json'), 'w') as file:
            file.write(json.dumps(vars(args), sort_keys=True, separators=('\n', ':')))

        with open(os.path.join(args.out_fold, 'train_loss.log'), 'w') as file:
            file.write("Start recording training loss ...\n")
        with open(os.path.join(args.out_fold, 'dev_loss.log'), 'w') as file:
            file.write("Start recording validation loss ...\n")
        with open(os.path.join(args.out_fold, 'test_loss.log'), 'w') as file:
            file.write("Start recording test loss ...\n")

    args.cuda = torch.cuda.is_available()
    print('Cuda device available: ', args.cuda)
    args.device = torch.device("cuda" if args.cuda else "cpu")

    return args

def adjust_learning_rate(args, lr, optimizer, epoch_num):
    lr = lr * (args.lr_decay ** (epoch_num // args.interval))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

def adjust_lambda_(args, epoch_num):
    args.lambda_ = 2 / (1 + np.exp(- 0.001 * epoch_num)) - 1 + 1e-9

def shuffle(feat, tags, labels):
    shuffle_index = torch.randperm(labels.shape[0])
    feat = feat[shuffle_index]
    tags = tags[shuffle_index]
    labels = labels[shuffle_index]
    # this_len = this_len[shuffle_index]
    return feat, tags, labels

def train(args):
    torch.set_default_tensor_type(torch.FloatTensor)

    # initialize model
    if args.model == 'resnet':
        feat_model = ResNet(11, args.enc_dim, resnet_type='18', nclasses=2).to(args.device)
    elif args.model == 'lcnn':
        feat_model = LCNN(4, args, nclasses=2).to(args.device)
    elif args.model == 'rawnet':
        assert args.feat == "Raw"
        with open("./model_config_RawNet.yml", 'r') as f_yaml:
            parser1 = yaml.safe_load(f_yaml)
            feat_model = RawNet(parser1["model"], args).to(args.device)

    if args.continue_training:
        feat_model = torch.load(os.path.join(args.out_fold, 'anti-spoofing_feat_model.pt')).to(args.device)
    feat_optimizer = torch.optim.Adam(feat_model.parameters(), lr=args.lr,
                                      betas=(args.beta_1, args.beta_2), eps=args.eps, weight_decay=0.0005)

    training_set = ASVspoof2019LASpec(args.path_to_database, args.path_to_features, 'train',
                                    args.feat, feat_len=args.feat_len)
    validation_set = ASVspoof2019LASpec(args.path_to_database, args.path_to_features, 'dev',
                                    args.feat, feat_len=args.feat_len)

    trainDataLoader = DataLoader(training_set, batch_size=args.batch_size,
                                 shuffle=True, num_workers=args.num_workers, collate_fn=training_set.collate_fn)
    valDataLoader = DataLoader(validation_set, batch_size=args.batch_size,
                               shuffle=True, num_workers=args.num_workers, collate_fn=validation_set.collate_fn)

    test_set = ASVspoof2019LASpec(args.path_to_database, args.path_to_features, "eval", args.feat,
                            feat_len=args.feat_len)
    testDataLoader = DataLoader(test_set, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, collate_fn=test_set.collate_fn)

    feat, _, _, _, _ = training_set[23]
    print("Feature shape", feat.shape)

    criterion = nn.CrossEntropyLoss().to(args.device)

    if args.loss == "ocsoftmax":
        ocsoftmax = OCSoftmax(args.enc_dim, m_real=args.m_real, m_fake=args.m_fake, alpha=args.alpha).to(args.device)
        ocsoftmax.train()
        ocsoftmax_optimizer = torch.optim.SGD(ocsoftmax.parameters(), lr=args.lr)
    elif args.loss == "isolate":
        iso_loss = IsolateLoss(2, args.enc_dim, r_real=args.r_real, r_fake=args.r_fake).to(args.device)
        iso_loss.train()
        iso_optimizer = torch.optim.SGD(iso_loss.parameters(), lr=args.lr)
    elif args.loss == "scl":
        scl_loss = SingleCenterLoss(2, args.enc_dim, m=args.scale_factor).to(args.device)
        scl_loss.train()
        scl_optimizer = torch.optim.SGD(scl_loss.parameters(), lr=args.lr)
    elif args.loss == "amsoftmax":
        amsoftmax_loss = AMSoftmax(2, args.enc_dim, s=args.alpha, m=args.m_real).to(args.device)
        amsoftmax_loss.train()
        amsoftmax_optimizer = torch.optim.SGD(amsoftmax_loss.parameters(), lr=0.01)
    elif args.loss == "angulariso":
        angulariso = AngularIsoLoss(args.enc_dim, m_real=args.m_real, m_fake=args.m_fake, alpha=args.alpha).to(args.device)
        angulariso.train()
        angulariso_optimizer = torch.optim.SGD(angulariso.parameters(), lr=args.lr)

    early_stop_cnt = 0
    prev_loss = 1e8

    monitor_loss = args.loss

    for epoch_num in tqdm(range(args.num_epochs)):
        genuine_feats, ip1_loader, tag_loader, idx_loader = [], [], [], []
        feat_model.train()
        trainlossDict = defaultdict(list)
        devlossDict = defaultdict(list)
        testlossDict = defaultdict(list)
        adjust_learning_rate(args, args.lr, feat_optimizer, epoch_num)
        if args.loss == "ocsoftmax":
            adjust_learning_rate(args, args.lr, ocsoftmax_optimizer, epoch_num)
        elif args.loss == "isolate":
            adjust_learning_rate(args, args.lr, iso_optimizer, epoch_num)
        elif args.loss == "scl":
            adjust_learning_rate(args, args.lr, scl_optimizer, epoch_num)
        elif args.loss == "amsoftmax":
            adjust_learning_rate(args, args.lr, amsoftmax_optimizer, epoch_num)
        elif args.loss == "angulariso":
            adjust_learning_rate(args, args.lr, angulariso_optimizer, epoch_num)
        adjust_lambda_(args, epoch_num)
        print('\nEpoch: %d ' % (epoch_num + 1))
        correct_m, total_m, correct_c, total_c, correct_v, total_v = 0, 0, 0, 0, 0, 0

        for i, (feat, audio_fn, tags, labels, channel) in enumerate(tqdm(trainDataLoader)):
            ## data prep
            if args.feat == "Raw":
                feat = feat.to(args.device)
            else:
                feat = feat.transpose(2,3).to(args.device)
            tags = tags.to(args.device)
            labels = labels.to(args.device)

            # Train the embedding network
            ## forward
            feats, feat_outputs = feat_model(feat)

            ## loss calculate
            if args.loss == "softmax":
                feat_loss = criterion(feat_outputs, labels)
            elif args.loss == "ocsoftmax":
                ocsoftmaxloss, _ = ocsoftmax(feats, labels)
                feat_loss = ocsoftmaxloss
            elif args.loss == "isolate":
                isoloss, _ = iso_loss(feats, labels)
                feat_loss = isoloss
            elif args.loss == "scl":
                sclloss, _ = scl_loss(feats, labels)
                feat_loss = criterion(feat_outputs, labels) + sclloss * args.weight_loss
            elif args.loss == "amsoftmax":
                outputs, moutputs = amsoftmax_loss(feats, labels)
                feat_loss = criterion(moutputs, labels)
            elif args.loss == "angulariso":
                angularisoloss, _ = angulariso(feats, labels)
                feat_loss = angularisoloss

            ## backward
            if args.loss == "softmax":
                trainlossDict[args.loss].append(feat_loss.item())
                feat_optimizer.zero_grad()
                feat_loss.backward()
                feat_optimizer.step()
            elif args.loss == "ocsoftmax":
                ocsoftmax_optimizer.zero_grad()
                trainlossDict[args.loss].append(ocsoftmaxloss.item())
                feat_optimizer.zero_grad()
                feat_loss.backward()
                feat_optimizer.step()
                ocsoftmax_optimizer.step()
            elif args.loss == "isolate":
                iso_optimizer.zero_grad()
                trainlossDict[args.loss].append(isoloss.item())
                feat_optimizer.zero_grad()
                feat_loss.backward()
                feat_optimizer.step()
                iso_optimizer.step()
            elif args.loss == "scl":
                scl_optimizer.zero_grad()
                trainlossDict[args.loss].append(sclloss.item())
                feat_optimizer.zero_grad()
                feat_loss.backward()
                feat_optimizer.step()
                scl_optimizer.step()
            elif args.loss == "amsoftmax":
                trainlossDict[args.loss].append(feat_loss.item())
                feat_optimizer.zero_grad()
                amsoftmax_optimizer.zero_grad()
                feat_loss.backward()
                feat_optimizer.step()
                amsoftmax_optimizer.step()
            elif args.loss == "angulariso":
                angulariso_optimizer.zero_grad()
                trainlossDict[args.loss].append(angularisoloss.item())
                feat_optimizer.zero_grad()
                feat_loss.backward()
                feat_optimizer.step()
                angulariso_optimizer.step()

            ## record
            ip1_loader.append(feats)
            idx_loader.append((labels))
            tag_loader.append((tags))


            with open(os.path.join(args.out_fold, "train_loss.log"), "a") as log:
                log.write(str(epoch_num) + "\t" + str(i) + "\t" +
                            str(trainlossDict[monitor_loss][-1]) + "\n")


        # Val the model
        # eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance)
        feat_model.eval()
        with torch.no_grad():
            ip1_loader, tag_loader, idx_loader, score_loader = [], [], [], []
            for i, (feat, audio_fn, tags, labels, channel) in enumerate(tqdm(valDataLoader)):
                if args.feat == "Raw":
                    feat = feat.to(args.device)
                else:
                    feat = feat.transpose(2, 3).to(args.device)

                tags = tags.to(args.device)
                labels = labels.to(args.device)
                feat, tags, labels = shuffle(feat, tags, labels)
                feats, feat_outputs = feat_model(feat)

                if args.loss == "softmax":
                    feat_loss = criterion(feat_outputs, labels)
                    score = F.softmax(feat_outputs, dim=1)[:, 0]
                    devlossDict[args.loss].append(feat_loss.item())
                elif args.loss == "ocsoftmax":
                    ocsoftmaxloss, score = ocsoftmax(feats, labels)
                    devlossDict[args.loss].append(ocsoftmaxloss.item())
                elif args.loss == "isolate":
                    isoloss, score = iso_loss(feats, labels)
                    devlossDict[args.loss].append(isoloss.item())
                elif args.loss == "scl":
                    sclloss, score = scl_loss(feats, labels)
                    sclloss = criterion(feat_outputs, labels) + sclloss * args.weight_loss
                    devlossDict[args.loss].append(sclloss.item())
                elif args.loss == "amsoftmax":
                    outputs, moutputs = amsoftmax_loss(feats, labels)
                    feat_loss = criterion(moutputs, labels)
                    score = F.softmax(outputs, dim=1)[:, 0]
                    devlossDict[args.loss].append(feat_loss.item())
                elif args.loss == "angulariso":
                    angularisoloss, score = angulariso(feats, labels)
                    devlossDict[args.loss].append(angularisoloss.item())

                ip1_loader.append(feats)
                idx_loader.append((labels))
                tag_loader.append((tags))

                score_loader.append(score)

            scores = torch.cat(score_loader, 0).data.cpu().numpy()
            labels = torch.cat(idx_loader, 0).data.cpu().numpy()
            eer = em.compute_eer(scores[labels == 0], scores[labels == 1])[0]

            with open(os.path.join(args.out_fold, "dev_loss.log"), "a") as log:
                log.write(str(epoch_num) + "\t" +
                            str(np.nanmean(devlossDict[monitor_loss])) + "\t" +
                            str(eer) +"\n")
            print("Val EER: {}".format(eer))


        if args.test_on_eval:
            if (epoch_num + 1) % args.test_interval == 0:
                with torch.no_grad():
                    ip1_loader, tag_loader, idx_loader, score_loader = [], [], [], []
                    for i, (feat, audio_fn, tags, labels, channel) in enumerate(tqdm(testDataLoader)):
                        if args.feat == "Raw":
                            feat = feat.to(args.device)
                        else:
                            feat = feat.transpose(2, 3).to(args.device)
                        tags = tags.to(args.device)
                        labels = labels.to(args.device)
                        feats, feat_outputs = feat_model(feat)
                        if args.loss == "softmax":
                            feat_loss = criterion(feat_outputs, labels)
                            score = F.softmax(feat_outputs, dim=1)[:, 0]
                            testlossDict[args.loss].append(feat_loss.item())
                        elif args.loss == "ocsoftmax":
                            ocsoftmaxloss, score = ocsoftmax(feats, labels)
                            testlossDict[args.loss].append(ocsoftmaxloss.item())
                        elif args.loss == "isolate":
                            isoloss, score = iso_loss(feats, labels)
                            testlossDict[args.loss].append(isoloss.item())
                        elif args.loss == "scl":
                            sclloss, score = scl_loss(feats, labels)
                            testlossDict[args.loss].append(sclloss.item())
                        elif args.loss == "amsoftmax":
                            outputs, moutputs = amsoftmax_loss(feats, labels)
                            feat_loss = criterion(moutputs, labels)
                            score = F.softmax(outputs, dim=1)[:, 0]
                            testlossDict[args.loss].append(feat_loss.item())
                        elif args.loss == "angulariso":
                            angularisoloss, score = angulariso(feats, labels)
                            testlossDict[args.loss].append(angularisoloss.item())

                        ip1_loader.append(feats)
                        idx_loader.append((labels))
                        tag_loader.append((tags))

                        score_loader.append(score)

                    scores = torch.cat(score_loader, 0).data.cpu().numpy()
                    labels = torch.cat(idx_loader, 0).data.cpu().numpy()
                    eer = em.compute_eer(scores[labels == 0], scores[labels == 1])[0]

                    with open(os.path.join(args.out_fold, "test_loss.log"), "a") as log:
                        log.write(str(epoch_num) + "\t" + str(np.nanmean(testlossDict[monitor_loss])) + "\t" + str(eer) + "\n")
                    print("Test EER: {}".format(eer))


        valLoss = np.nanmean(devlossDict[monitor_loss])
        if (epoch_num + 1) % args.save_interval == 0:
            # Save the model checkpoint
            torch.save(feat_model, os.path.join(args.out_fold, 'checkpoint',
                                                'anti-spoofing_feat_model_%d.pt' % (epoch_num + 1)))
            if args.loss == "ocsoftmax":
                loss_model = ocsoftmax
            elif args.loss == "isolate":
                loss_model = iso_loss
            elif args.loss == "scl":
                loss_model = scl_loss
            elif args.loss == "softmax":
                loss_model = None
            elif args.loss == "amsoftmax":
                loss_model = amsoftmax_loss
            elif args.loss == "angulariso":
                loss_model = angulariso
            else:
                print("What is your loss? You may encounter error.")
            torch.save(loss_model, os.path.join(args.out_fold, 'checkpoint',
                                                'anti-spoofing_loss_model_%d.pt' % (epoch_num + 1)))

        if valLoss < prev_loss:
            torch.save(feat_model, os.path.join(args.out_fold, 'anti-spoofing_feat_model.pt'))
            if args.loss == "ocsoftmax":
                loss_model = ocsoftmax
            elif args.loss == "isolate":
                loss_model = iso_loss
            elif args.loss == "scl":
                loss_model = scl_loss
            elif args.loss == "softmax":
                loss_model = None
            elif args.loss == "amsoftmax":
                loss_model = amsoftmax_loss
            elif args.loss == "angulariso":
                loss_model = angulariso
            else:
                print("What is your loss? You may encounter error.")
            torch.save(loss_model, os.path.join(args.out_fold, 'anti-spoofing_loss_model.pt'))

            prev_loss = valLoss
            early_stop_cnt = 0
        else:
            early_stop_cnt += 1

        if early_stop_cnt == 50:
            with open(os.path.join(args.out_fold, 'args.json'), 'a') as res_file:
                res_file.write('\nTrained Epochs: %d\n' % (epoch_num - 49))
            break


if __name__ == "__main__":
    args = initParams()
    train(args)


models/feat+resnet/train.py
yongyizang/SingFake
import torch
import argparse
import os
import json
import shutil
import numpy as np
from model import *
from dataset import *
from torch.utils.data import DataLoader
from loss import *
from collections import defaultdict
from tqdm import tqdm
from utils import str2bool, setup_seed
import eval_metrics as em
import yaml

torch.set_default_tensor_type(torch.FloatTensor)

def initParams():
    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument('--seed', type=int, help="random number seed", default=1000)

    # Data folder prepare
    parser.add_argument("-d", "--path_to_database", type=str, help="dataset path",
                        default='/data/neil/DS_10283_3336/')
    parser.add_argument("-f", "--path_to_features", type=str, help="features path",
                        default='/data2/neil/ASVspoof2019LA/')
    parser.add_argument("-p", "--path_to_protocol", type=str, help="protocol path",
                        default='/data/neil/DS_10283_3336/LA/ASVspoof2019_LA_cm_protocols/')
    parser.add_argument("-o", "--out_fold", type=str, help="output folder", required=True, default='./models/try/')

    # Dataset prepare
    parser.add_argument("--feat", type=str, help="which feature to use", default='LFCC',
                        choices=["CQCC", "LFCC", "Raw", "MFCC"])
    parser.add_argument("--feat_len", type=int, help="features length", default=500)
    parser.add_argument("--enc_dim", type=int, help="encoding dimension", default=256)

    parser.add_argument('-m', '--model', help='Model arch', default='resnet',
                        choices=['resnet', 'lcnn', 'rawnet'])

    # Training hyperparameters
    parser.add_argument('--num_epochs', type=int, default=100, help="Number of epochs for training")
    parser.add_argument('--batch_size', type=int, default=128, help="Mini batch size for training")
    parser.add_argument('--lr', type=float, default=0.0003, help="learning rate")
    parser.add_argument('--lr_decay', type=float, default=0.5, help="decay learning rate")
    parser.add_argument('--interval', type=int, default=20, help="interval to decay lr")

    parser.add_argument('--beta_1', type=float, default=0.9, help="bata_1 for Adam")
    parser.add_argument('--beta_2', type=float, default=0.999, help="beta_2 for Adam")
    parser.add_argument('--eps', type=float, default=1e-8, help="epsilon for Adam")
    parser.add_argument("--gpu", type=str, help="GPU index", default="1")
    parser.add_argument('--num_workers', type=int, default=0, help="number of workers")

    parser.add_argument('-l', '--loss', type=str, default="ocsoftmax",
                        choices=["softmax", "amsoftmax", "ocsoftmax", "isolate", "scl", "angulariso"], help="loss for training")
    parser.add_argument('--weight_loss', type=float, default=0.5, help="weight for other loss")
    parser.add_argument('--m_real', type=float, default=0.5, help="m_real for ocsoftmax loss")
    parser.add_argument('--m_fake', type=float, default=0.2, help="m_fake for ocsoftmax loss")
    parser.add_argument('--r_real', type=float, default=25.0, help="r_real for isolate loss")
    parser.add_argument('--r_fake', type=float, default=75.0, help="r_fake for isolate loss")
    parser.add_argument('--alpha', type=float, default=20, help="scale factor for amsoftmax and ocsoftmax loss")
    parser.add_argument('--scale_factor', type=float, default=0.5, help="scale factor for single center loss")

    parser.add_argument('--continue_training', action='store_true', help="continue training with trained model")

    parser.add_argument('--is_mixture', type=str2bool, nargs='?', const=True, default=False,
                        help="whether use mixture or vocals in training")
    parser.add_argument('--test_on_eval', action='store_true',
                        help="whether to run EER on the evaluation set")
    parser.add_argument('--test_interval', type=int, default=5, help="test on eval for every how many epochs")
    parser.add_argument('--save_interval', type=int, default=5, help="save checkpoint model for every how many epochs")

    args = parser.parse_args()

    # Change this to specify GPU
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu

    # Set seeds
    setup_seed(args.seed)

    if any([args.feat == "Raw", args.model == "rawnet", args.feat_len > 16000]):
        assert all([args.feat == "Raw", args.model == "rawnet", args.feat_len > 16000])

    if args.continue_training:
        pass
    else:
        # Path for output data
        if not os.path.exists(args.out_fold):
            os.makedirs(args.out_fold)
        else:
            shutil.rmtree(args.out_fold)
            os.mkdir(args.out_fold)

        # Folder for intermediate results
        if not os.path.exists(os.path.join(args.out_fold, 'checkpoint')):
            os.makedirs(os.path.join(args.out_fold, 'checkpoint'))
        else:
            shutil.rmtree(os.path.join(args.out_fold, 'checkpoint'))
            os.mkdir(os.path.join(args.out_fold, 'checkpoint'))

        # Path for input data
        # assert os.path.exists(args.path_to_database)
        # assert os.path.exists(args.path_to_features)

        # Save training arguments
        with open(os.path.join(args.out_fold, 'args.json'), 'w') as file:
            file.write(json.dumps(vars(args), sort_keys=True, separators=('\n', ':')))

        with open(os.path.join(args.out_fold, 'train_loss.log'), 'w') as file:
            file.write("Start recording training loss ...\n")
        with open(os.path.join(args.out_fold, 'dev_loss.log'), 'w') as file:
            file.write("Start recording validation loss ...\n")
        with open(os.path.join(args.out_fold, 'test_loss.log'), 'w') as file:
            file.write("Start recording test loss ...\n")

    args.cuda = torch.cuda.is_available()
    print('Cuda device available: ', args.cuda)
    args.device = torch.device("cuda" if args.cuda else "cpu")

    return args

def adjust_learning_rate(args, lr, optimizer, epoch_num):
    lr = lr * (args.lr_decay ** (epoch_num // args.interval))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

def adjust_lambda_(args, epoch_num):
    args.lambda_ = 2 / (1 + np.exp(- 0.001 * epoch_num)) - 1 + 1e-9

def shuffle(feat, tags, labels):
    shuffle_index = torch.randperm(labels.shape[0])
    feat = feat[shuffle_index]
    tags = tags[shuffle_index]
    labels = labels[shuffle_index]
    # this_len = this_len[shuffle_index]
    return feat, tags, labels

def train(args):
    torch.set_default_tensor_type(torch.FloatTensor)

    # initialize model
    if args.model == 'resnet':
        feat_model = ResNet(3, args.enc_dim, resnet_type='18', nclasses=2).to(args.device)
    elif args.model == 'lcnn':
        feat_model = LCNN(4, args, nclasses=2).to(args.device)
    elif args.model == 'rawnet':
        assert args.feat == "Raw"
        with open("./model_config_RawNet.yml", 'r') as f_yaml:
            parser1 = yaml.safe_load(f_yaml)
            feat_model = RawNet(parser1["model"], args).to(args.device)

    if args.continue_training:
        feat_model = torch.load(os.path.join(args.out_fold, 'anti-spoofing_feat_model.pt')).to(args.device)
    feat_optimizer = torch.optim.Adam(feat_model.parameters(), lr=args.lr,
                                      betas=(args.beta_1, args.beta_2), eps=args.eps, weight_decay=0.0005)

    training_set = Dataset_SingFake(os.path.join(args.path_to_database, "train"), args.is_mixture)
    validation_set = Dataset_SingFake(os.path.join(args.path_to_database, "dev"), args.is_mixture)

    trainDataLoader = DataLoader(training_set, batch_size=args.batch_size,
                                 shuffle=True, num_workers=args.num_workers)
    valDataLoader = DataLoader(validation_set, batch_size=args.batch_size,
                               shuffle=True, num_workers=args.num_workers)

    test_set = Dataset_SingFake(os.path.join(args.path_to_database, "test"), args.is_mixture)
    testDataLoader = DataLoader(test_set, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers)

    feat, _ = training_set[23]
    print("Feature shape", feat.shape)

    criterion = nn.CrossEntropyLoss().to(args.device)

    if args.loss == "ocsoftmax":
        ocsoftmax = OCSoftmax(args.enc_dim, m_real=args.m_real, m_fake=args.m_fake, alpha=args.alpha).to(args.device)
        ocsoftmax.train()
        ocsoftmax_optimizer = torch.optim.SGD(ocsoftmax.parameters(), lr=args.lr)
    elif args.loss == "isolate":
        iso_loss = IsolateLoss(2, args.enc_dim, r_real=args.r_real, r_fake=args.r_fake).to(args.device)
        iso_loss.train()
        iso_optimizer = torch.optim.SGD(iso_loss.parameters(), lr=args.lr)
    elif args.loss == "scl":
        scl_loss = SingleCenterLoss(2, args.enc_dim, m=args.scale_factor).to(args.device)
        scl_loss.train()
        scl_optimizer = torch.optim.SGD(scl_loss.parameters(), lr=args.lr)
    elif args.loss == "amsoftmax":
        amsoftmax_loss = AMSoftmax(2, args.enc_dim, s=args.alpha, m=args.m_real).to(args.device)
        amsoftmax_loss.train()
        amsoftmax_optimizer = torch.optim.SGD(amsoftmax_loss.parameters(), lr=0.01)
    elif args.loss == "angulariso":
        angulariso = AngularIsoLoss(args.enc_dim, m_real=args.m_real, m_fake=args.m_fake, alpha=args.alpha).to(args.device)
        angulariso.train()
        angulariso_optimizer = torch.optim.SGD(angulariso.parameters(), lr=args.lr)

    early_stop_cnt = 0
    prev_loss = 1e8

    monitor_loss = args.loss

    for epoch_num in tqdm(range(args.num_epochs)):
        genuine_feats, ip1_loader, tag_loader, idx_loader = [], [], [], []
        feat_model.train()
        trainlossDict = defaultdict(list)
        devlossDict = defaultdict(list)
        testlossDict = defaultdict(list)
        adjust_learning_rate(args, args.lr, feat_optimizer, epoch_num)
        if args.loss == "ocsoftmax":
            adjust_learning_rate(args, args.lr, ocsoftmax_optimizer, epoch_num)
        elif args.loss == "isolate":
            adjust_learning_rate(args, args.lr, iso_optimizer, epoch_num)
        elif args.loss == "scl":
            adjust_learning_rate(args, args.lr, scl_optimizer, epoch_num)
        elif args.loss == "amsoftmax":
            adjust_learning_rate(args, args.lr, amsoftmax_optimizer, epoch_num)
        elif args.loss == "angulariso":
            adjust_learning_rate(args, args.lr, angulariso_optimizer, epoch_num)
        print('\nEpoch: %d ' % (epoch_num + 1))
        correct_m, total_m, correct_c, total_c, correct_v, total_v = 0, 0, 0, 0, 0, 0

        for i, (feat, labels) in enumerate(tqdm(trainDataLoader)):
            # if i > 2: break # debug purpose
            ## data prep
            if args.feat == "Raw":
                feat = feat.to(args.device)
            else:
                feat = feat.transpose(2,3).to(args.device)
            labels = labels.to(args.device)

            # Train the embedding network
            ## forward
            feats, feat_outputs = feat_model(feat)

            ## loss calculate
            if args.loss == "softmax":
                feat_loss = criterion(feat_outputs, labels)
            elif args.loss == "ocsoftmax":
                ocsoftmaxloss, _ = ocsoftmax(feats, labels)
                feat_loss = ocsoftmaxloss
            elif args.loss == "isolate":
                isoloss, _ = iso_loss(feats, labels)
                feat_loss = isoloss
            elif args.loss == "scl":
                sclloss, _ = scl_loss(feats, labels)
                feat_loss = criterion(feat_outputs, labels) + sclloss * args.weight_loss
            elif args.loss == "amsoftmax":
                outputs, moutputs = amsoftmax_loss(feats, labels)
                feat_loss = criterion(moutputs, labels)
            elif args.loss == "angulariso":
                angularisoloss, _ = angulariso(feats, labels)
                feat_loss = angularisoloss


            ## backward
            if args.loss == "softmax":
                trainlossDict[args.loss].append(feat_loss.item())
                feat_optimizer.zero_grad()
                feat_loss.backward()
                feat_optimizer.step()
            elif args.loss == "ocsoftmax":
                ocsoftmax_optimizer.zero_grad()
                trainlossDict[args.loss].append(ocsoftmaxloss.item())
                feat_optimizer.zero_grad()
                feat_loss.backward()
                feat_optimizer.step()
                ocsoftmax_optimizer.step()
            elif args.loss == "isolate":
                iso_optimizer.zero_grad()
                trainlossDict[args.loss].append(isoloss.item())
                feat_optimizer.zero_grad()
                feat_loss.backward()
                feat_optimizer.step()
                iso_optimizer.step()
            elif args.loss == "scl":
                scl_optimizer.zero_grad()
                trainlossDict[args.loss].append(sclloss.item())
                feat_optimizer.zero_grad()
                feat_loss.backward()
                feat_optimizer.step()
                scl_optimizer.step()
            elif args.loss == "amsoftmax":
                trainlossDict[args.loss].append(feat_loss.item())
                feat_optimizer.zero_grad()
                amsoftmax_optimizer.zero_grad()
                feat_loss.backward()
                feat_optimizer.step()
                amsoftmax_optimizer.step()
            elif args.loss == "angulariso":
                angulariso_optimizer.zero_grad()
                trainlossDict[args.loss].append(angularisoloss.item())
                feat_optimizer.zero_grad()
                feat_loss.backward()
                feat_optimizer.step()
                angulariso_optimizer.step()


            ## record
            ip1_loader.append(feats)
            idx_loader.append((labels))

            with open(os.path.join(args.out_fold, "train_loss.log"), "a") as log:
                log.write(str(epoch_num) + "\t" + str(i) + "\t" +
                            str(trainlossDict[monitor_loss][-1]) + "\n")


        # Val the model
        # eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance)
        feat_model.eval()
        with torch.no_grad():
            ip1_loader, tag_loader, idx_loader, score_loader = [], [], [], []
            for i, (feat, labels) in enumerate(tqdm(valDataLoader)):
                if args.feat == "Raw":
                    feat = feat.to(args.device)
                else:
                    feat = feat.transpose(2, 3).to(args.device)

                labels = labels.to(args.device)
                feats, feat_outputs = feat_model(feat)

                if args.loss == "softmax":
                    feat_loss = criterion(feat_outputs, labels)
                    score = F.softmax(feat_outputs, dim=1)[:, 0]
                    devlossDict[args.loss].append(feat_loss.item())
                elif args.loss == "ocsoftmax":
                    ocsoftmaxloss, score = ocsoftmax(feats, labels)
                    devlossDict[args.loss].append(ocsoftmaxloss.item())
                elif args.loss == "isolate":
                    isoloss, score = iso_loss(feats, labels)
                    devlossDict[args.loss].append(isoloss.item())
                elif args.loss == "scl":
                    sclloss, score = scl_loss(feats, labels)
                    sclloss = criterion(feat_outputs, labels) + sclloss * args.weight_loss
                    devlossDict[args.loss].append(sclloss.item())
                elif args.loss == "amsoftmax":
                    outputs, moutputs = amsoftmax_loss(feats, labels)
                    feat_loss = criterion(moutputs, labels)
                    score = F.softmax(outputs, dim=1)[:, 0]
                    devlossDict[args.loss].append(feat_loss.item())
                elif args.loss == "angulariso":
                    angularisoloss, score = angulariso(feats, labels)
                    devlossDict[args.loss].append(angularisoloss.item())


                ip1_loader.append(feats)
                idx_loader.append((labels))

                score_loader.append(score)

            scores = torch.cat(score_loader, 0).data.cpu().numpy()
            labels = torch.cat(idx_loader, 0).data.cpu().numpy()
            eer = em.compute_eer(scores[labels == 0], scores[labels == 1])[0]

            with open(os.path.join(args.out_fold, "dev_loss.log"), "a") as log:
                log.write(str(epoch_num) + "\t" +
                            str(np.nanmean(devlossDict[monitor_loss])) + "\t" +
                            str(eer) +"\n")
            print("Val EER: {}".format(eer))


        if args.test_on_eval:
            if (epoch_num + 1) % args.test_interval == 0:
                with torch.no_grad():
                    ip1_loader, tag_loader, idx_loader, score_loader = [], [], [], []
                    for i, (feat, labels) in enumerate(tqdm(testDataLoader)):
                        if args.feat == "Raw":
                            feat = feat.to(args.device)
                        else:
                            feat = feat.transpose(2, 3).to(args.device)

                        labels = labels.to(args.device)
                        feats, feat_outputs = feat_model(feat)
                        if args.loss == "softmax":
                            feat_loss = criterion(feat_outputs, labels)
                            score = F.softmax(feat_outputs, dim=1)[:, 0]
                            testlossDict[args.loss].append(feat_loss.item())
                        elif args.loss == "ocsoftmax":
                            ocsoftmaxloss, score = ocsoftmax(feats, labels)
                            testlossDict[args.loss].append(ocsoftmaxloss.item())
                        elif args.loss == "isolate":
                            isoloss, score = iso_loss(feats, labels)
                            testlossDict[args.loss].append(isoloss.item())
                        elif args.loss == "scl":
                            sclloss, score = scl_loss(feats, labels)
                            testlossDict[args.loss].append(sclloss.item())
                        elif args.loss == "amsoftmax":
                            outputs, moutputs = amsoftmax_loss(feats, labels)
                            feat_loss = criterion(moutputs, labels)
                            score = F.softmax(outputs, dim=1)[:, 0]
                            testlossDict[args.loss].append(feat_loss.item())
                        elif args.loss == "angulariso":
                            angularisoloss, score = angulariso(feats, labels)
                            testlossDict[args.loss].append(angularisoloss.item())

                        ip1_loader.append(feats)
                        idx_loader.append((labels))

                        score_loader.append(score)

                    scores = torch.cat(score_loader, 0).data.cpu().numpy()
                    labels = torch.cat(idx_loader, 0).data.cpu().numpy()
                    eer = em.compute_eer(scores[labels == 0], scores[labels == 1])[0]

                    with open(os.path.join(args.out_fold, "test_loss.log"), "a") as log:
                        log.write(str(epoch_num) + "\t" + str(np.nanmean(testlossDict[monitor_loss])) + "\t" + str(eer) + "\n")
                    print("Test EER: {}".format(eer))


        valLoss = np.nanmean(devlossDict[monitor_loss])
        if (epoch_num + 1) % args.save_interval == 0:
            # Save the model checkpoint
            torch.save(feat_model, os.path.join(args.out_fold, 'checkpoint',
                                                'anti-spoofing_feat_model_%d.pt' % (epoch_num + 1)))
            if args.loss == "ocsoftmax":
                loss_model = ocsoftmax
            elif args.loss == "isolate":
                loss_model = iso_loss
            elif args.loss == "scl":
                loss_model = scl_loss
            elif args.loss == "softmax":
                loss_model = None
            elif args.loss == "amsoftmax":
                loss_model = amsoftmax_loss
            elif args.loss == "angulariso":
                loss_model = angulariso
            else:
                print("What is your loss? You may encounter error.")
            torch.save(loss_model, os.path.join(args.out_fold, 'checkpoint',
                                                'anti-spoofing_loss_model_%d.pt' % (epoch_num + 1)))

        if valLoss < prev_loss:
            torch.save(feat_model, os.path.join(args.out_fold, 'anti-spoofing_feat_model.pt'))
            if args.loss == "ocsoftmax":
                loss_model = ocsoftmax
            elif args.loss == "isolate":
                loss_model = iso_loss
            elif args.loss == "scl":
                loss_model = scl_loss
            elif args.loss == "softmax":
                loss_model = None
            elif args.loss == "amsoftmax":
                loss_model = amsoftmax_loss
            elif args.loss == "angulariso":
                loss_model = angulariso
            else:
                print("What is your loss? You may encounter error.")
            torch.save(loss_model, os.path.join(args.out_fold, 'anti-spoofing_loss_model.pt'))

            prev_loss = valLoss
            early_stop_cnt = 0
        else:
            early_stop_cnt += 1

        if early_stop_cnt == 50:
            with open(os.path.join(args.out_fold, 'args.json'), 'a') as res_file:
                res_file.write('\nTrained Epochs: %d\n' % (epoch_num - 49))
            break


if __name__ == "__main__":
    args = initParams()
    train(args)


models/feat+resnet/train_SVspoof.py
yongyizang/SingFake
import os
import torch
import random
import numpy as np
from distutils import util

def str2bool(v):
    return bool(util.strtobool(v))

def setup_seed(random_seed, cudnn_deterministic=True):
    """ set_random_seed(random_seed, cudnn_deterministic=True)

    Set the random_seed for numpy, python, and cudnn

    input
    -----
      random_seed: integer random seed
      cudnn_deterministic: for torch.backends.cudnn.deterministic

    Note: this default configuration may result in RuntimeError
    see https://pytorch.org/docs/stable/notes/randomness.html
    """

    # # initialization
    # torch.manual_seed(random_seed)
    random.seed(random_seed)
    np.random.seed(random_seed)
    os.environ['PYTHONHASHSEED'] = str(random_seed)

    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(random_seed)
        torch.backends.cudnn.deterministic = cudnn_deterministic
        torch.backends.cudnn.benchmark = False


models/feat+resnet/utils.py
yongyizang/SingFake
import numpy as np
import os
import pickle
import matplotlib.pyplot as plt
import torch
from torch.utils.data import DataLoader
from dataset import ASVspoof2019LA
from tqdm import tqdm
from sklearn.manifold import TSNE
from sklearn.decomposition import PCA


def visualize_dev_and_eval(dev_feat, dev_labels, dev_tags, eval_feat, eval_labels, eval_tags,
                           center, seed, out_fold):
    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(8, 8))
    c = ['#7030a0', '#ff0000', '#ffff00']

    torch.manual_seed(888)
    num_centers, enc_dim = center.shape
    ind_dev = torch.randperm(dev_feat.shape[0])[:5000].numpy()
    ind_eval = torch.randperm(eval_feat.shape[0])[:5000].numpy()

    dev_feat_sample = dev_feat[ind_dev]
    eval_feat_sample = eval_feat[ind_eval]
    dev_lab_sam = dev_labels[ind_dev]
    eval_lab_sam = eval_labels[ind_eval]

    X = np.concatenate((center, dev_feat_sample, eval_feat_sample), axis=0)
    os.environ['PYTHONHASHSEED'] = str(668)
    np.random.seed(668)
    X_tsne = TSNE(random_state=seed, perplexity=40, early_exaggeration=40).fit_transform(X)
    center = X_tsne[:num_centers]
    feat_dev = X_tsne[num_centers:num_centers + 5000]
    feat_eval = X_tsne[num_centers + 5000:]
    pca = PCA(n_components=2)
    X_pca = pca.fit_transform(X)
    ex_ratio = pca.explained_variance_ratio_
    center_pca = X_pca[:num_centers]
    feat_pca_dev = X_pca[num_centers:num_centers + 5000]
    feat_pca_eval = X_pca[num_centers + 5000:]

    # t-SNE visualization
    ax1.plot(feat_dev[dev_lab_sam == 0, 0], feat_dev[dev_lab_sam == 0, 1], '.', c=c[0], markersize=1.2)
    ax1.plot(feat_dev[dev_lab_sam == 1, 0], feat_dev[dev_lab_sam == 1, 1], '.', c=c[1], markersize=1.2)
    ax1.axis('off')
    ax1.plot(center[:, 0], center[:, 1], 'x', c=c[2], markersize=5)
    plt.setp((ax2), xlim=ax1.get_xlim(), ylim=ax1.get_ylim())
    ax2.plot(center[:, 0], center[:, 1], 'x', c=c[2], markersize=5)

    ax2.plot(feat_eval[eval_lab_sam == 0, 0], feat_eval[eval_lab_sam == 0, 1], '.', c=c[0], markersize=1.2)
    ax2.plot(feat_eval[eval_lab_sam == 1, 0], feat_eval[eval_lab_sam == 1, 1], '.', c=c[1], markersize=1.2)
    ax2.axis('off')

    # PCA visualization
    ax3.plot(feat_pca_dev[dev_lab_sam == 0, 0], feat_pca_dev[dev_lab_sam == 0, 1], '.', c=c[0], markersize=1.2)
    ax3.plot(feat_pca_dev[dev_lab_sam == 1, 0], feat_pca_dev[dev_lab_sam == 1, 1], '.', c=c[1], markersize=1.2)
    ax3.axis('off')
    plt.setp((ax4), xlim=ax3.get_xlim(), ylim=ax3.get_ylim())
    ax4.plot(feat_pca_eval[eval_lab_sam == 0, 0], feat_pca_eval[eval_lab_sam == 0, 1], '.', c=c[0], markersize=1.2)
    ax4.plot(feat_pca_eval[eval_lab_sam == 1, 0], feat_pca_eval[eval_lab_sam == 1, 1], '.', c=c[1], markersize=1.2)
    ax4.axis('off')
    plt.savefig(os.path.join(out_fold, '_vis_feat.pdf'), dpi=500, bbox_inches="tight")
    plt.show()
    fig.clf()
    plt.close(fig)

def get_features(feat_model_path, part):
    model = torch.load(feat_model_path)
    dataset = ASVspoof2019LA('/data/neil/DS_10283_3336/',
                             '/data2/neil/ASVspoof2019LA/', part,
                             "LFCC", feat_len=500)
    dataLoader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=0,
                            collate_fn=dataset.collate_fn)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.eval()
    ip1_loader, tag_loader, idx_loader, score_loader = [], [], [], []
    for i, (feat, audio_fn, tags, labels, channel) in enumerate(tqdm(dataLoader)):
        feat = feat.transpose(2, 3).float().to(device)
        tags = tags.to(device)
        labels = labels.to(device)
        feats, _ = model(feat)
        ip1_loader.append(feats.detach().cpu().numpy())
        idx_loader.append((labels.detach().cpu().numpy()))
        tag_loader.append((tags.detach().cpu().numpy()))
    features = np.concatenate(ip1_loader, 0)
    labels = np.concatenate(idx_loader, 0)
    tags = np.concatenate(tag_loader, 0)
    gen_feats = features[labels == 0]
    return features, labels, tags


if __name__ == "__main__":
    os.environ["CUDA_VISIBLE_DEVICES"] = "3"
    device = torch.device("cuda")

    model_dir = "/data3/neil/hbas/models1106/resnet_ocsoftmax_500"

    feat_model_path = os.path.join(model_dir, "anti-spoofing_feat_model.pt")
    loss_model_path = os.path.join(model_dir, "anti-spoofing_loss_model.pt")

    center = torch.load(loss_model_path).center.detach().cpu().numpy()
    dev_feat, dev_labels, dev_tags = get_features(feat_model_path, "dev")
    eval_feat, eval_labels, eval_tags = get_features(feat_model_path, "eval")

    visualize_dev_and_eval(dev_feat, dev_labels, dev_tags, eval_feat, eval_labels, eval_tags, center, 88, model_dir)


models/feat+resnet/visualize.py
yongyizang/SingFake
import numpy as np
import soundfile as sf
import torch, os
from torch import Tensor
from torch.utils.data import Dataset
import librosa

___author__ = "Hemlata Tak, Jee-weon Jung"
__email__ = "tak@eurecom.fr, jeeweon.jung@navercorp.com"


def genSpoof_list(dir_meta, is_train=False, is_eval=False):
    d_meta = {}
    file_list = []
    with open(dir_meta, "r") as f:
        l_meta = f.readlines()

    if is_train:
        for line in l_meta:
            _, key, _, _, label = line.strip().split(" ")
            file_list.append(key)
            d_meta[key] = 1 if label == "bonafide" else 0
        return d_meta, file_list

    elif is_eval:
        for line in l_meta:
            _, key, _, _, _ = line.strip().split(" ")
            #key = line.strip()
            file_list.append(key)
        return file_list
    else:
        for line in l_meta:
            _, key, _, _, label = line.strip().split(" ")
            file_list.append(key)
            d_meta[key] = 1 if label == "bonafide" else 0
        return d_meta, file_list


def pad(x, max_len=64600):
    x_len = x.shape[0]
    if x_len >= max_len:
        return x[:max_len]
    # need to pad
    num_repeats = int(max_len / x_len) + 1
    padded_x = np.tile(x, (1, num_repeats))[:, :max_len][0]
    return padded_x


def pad_random(x: np.ndarray, max_len: int = 64600):
    x_len = x.shape[0]
    # if duration is already long enough
    if x_len >= max_len:
        stt = np.random.randint(x_len - max_len)
        return x[stt:stt + max_len]

    # if too short
    num_repeats = int(max_len / x_len) + 1
    padded_x = np.tile(x, (num_repeats))[:max_len]
    return padded_x


class Dataset_ASVspoof2019_train(Dataset):
    def __init__(self, list_IDs, labels, base_dir):
        """self.list_IDs	: list of strings (each string: utt key),
           self.labels      : dictionary (key: utt key, value: label integer)"""
        self.list_IDs = list_IDs
        self.labels = labels
        self.base_dir = base_dir
        self.cut = 64600  # take ~4 sec audio (64600 samples)

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

    def __getitem__(self, index):
        key = self.list_IDs[index]
        X, _ = sf.read(str(self.base_dir / f"flac/{key}.flac"))
        X_pad = pad_random(X, self.cut)
        x_inp = Tensor(X_pad)
        y = self.labels[key]
        return x_inp, y


class Dataset_ASVspoof2019_devNeval(Dataset):
    def __init__(self, list_IDs, base_dir):
        """self.list_IDs	: list of strings (each string: utt key),
        """
        self.list_IDs = list_IDs
        self.base_dir = base_dir
        self.cut = 64600  # take ~4 sec audio (64600 samples)

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

    def __getitem__(self, index):
        key = self.list_IDs[index]
        X, _ = sf.read(str(self.base_dir / f"flac/{key}.flac"))
        X_pad = pad(X, self.cut)
        x_inp = Tensor(X_pad)
        return x_inp, key

class Dataset_SingFake(Dataset):
    def __init__(self, base_dir, is_mixture=False, target_sr=16000):
        """
        base_dir should contain mixtures/ and vocals/ folders
        """
        self.base_dir = base_dir
        self.is_mixture = is_mixture
        self.target_sr = target_sr
        self.cut = 64600  # take ~4 sec audio (64600 samples)
        
        # get file list
        self.file_list = []
        if self.is_mixture:
            self.target_path = os.path.join(self.base_dir, "mixtures")
        else:
            self.target_path = os.path.join(self.base_dir, "vocals")
            
        print(self.target_path)
        
        assert os.path.exists(self.target_path), f"{self.target_path} does not exist!"
        
        for file in os.listdir(self.target_path):
            if file.endswith(".flac"):
                self.file_list.append(file[:-5])
    
    def __len__(self):
        return len(self.file_list)

    def __getitem__(self, index):
        key = self.file_list[index]
        file_path = os.path.join(self.target_path, key + ".flac")
        # X, _ = sf.read(file_path, samplerate=self.target_sr)
        try:
            X, _ = librosa.load(file_path, sr=self.target_sr, mono=False)
        except:
            return self.__getitem__(np.random.randint(len(self.file_list)))
        if X.shape[0] > 1:
            # if not mono, take random channel
            channel_id = np.random.randint(X.shape[0])
            X = X[channel_id]
            # X = np.expand_dims(X, axis=-1)
        X_pad = pad_random(X, self.cut)
        X_pad = X_pad / np.max(np.abs(X_pad))
        x_inp = Tensor(X_pad)
        y = int(key.split("_")[0])
        return x_inp, y
models/wav2vec2+AASIST/data_utils.py
yongyizang/SingFake
"""
Main script that trains, validates, and evaluates
various models including AASIST.

AASIST
Copyright (c) 2021-present NAVER Corp.
MIT license
"""
import argparse
import json
import os
import sys
import warnings
from importlib import import_module
from pathlib import Path
from shutil import copy
from typing import Dict, List, Union
from tqdm import tqdm
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchcontrib.optim import SWA
from data_utils import (Dataset_ASVspoof2019_train,
                        Dataset_ASVspoof2019_devNeval, genSpoof_list, Dataset_SingFake)
from evaluation import compute_eer
from utils import create_optimizer, seed_worker, set_seed, str_to_bool
from models.wav2vecAASIST import Wav2Vec2Model
warnings.filterwarnings("ignore", category=FutureWarning)

def main(args: argparse.Namespace) -> None:
    """
    Main function.
    Trains, validates, and evaluates the ASVspoof detection model.
    """
    # load experiment configurations
    with open(args.config, "r") as f_json:
        config = json.loads(f_json.read())
    model_config = config["model_config"]
    optim_config = config["optim_config"]
    optim_config["epochs"] = config["num_epochs"]
    track = config["track"]
    assert track in ["LA", "PA", "DF"], "Invalid track given"
    if "eval_all_best" not in config:
        config["eval_all_best"] = "True"
    if "freq_aug" not in config:
        config["freq_aug"] = "False"

    # make experiment reproducible
    set_seed(args.seed, config)

    # define database related paths
    output_dir = Path(args.output_dir)
    prefix_2019 = "ASVspoof2019.{}".format(track)
    database_path = Path(config["database_path"])
    dev_trial_path = (database_path /
                      "ASVspoof2019_{}_cm_protocols/{}.cm.dev.trl.txt".format(
                          track, prefix_2019))
    eval_trial_path = (
        database_path /
        "ASVspoof2019_{}_cm_protocols/{}.cm.eval.trl.txt".format(
            track, prefix_2019))

    # define model related paths
    model_tag = "{}_{}_ep{}_bs{}".format(
        track,
        os.path.splitext(os.path.basename(args.config))[0],
        config["num_epochs"], config["batch_size"])
    if args.comment:
        model_tag = model_tag + "_{}".format(args.comment)
    model_tag = output_dir / model_tag
    model_save_path = model_tag / "weights"
    eval_score_path = model_tag / config["eval_output"]
    writer = SummaryWriter(model_tag)
    os.makedirs(model_save_path, exist_ok=True)
    copy(args.config, model_tag / "config.conf")

    # set device
    gpu_id = args.gpu
    device = torch.device("cuda:{}".format(gpu_id) if torch.cuda.is_available() else "cpu")
    # device = "cuda" if torch.cuda.is_available() else "cpu"
    print("Device: {}".format(device))
    if device == "cpu":
        raise ValueError("GPU not detected!")

    # define model architecture
    # model = get_model(model_config, device)
    model = get_wav2vec2_model(model_config, device).to(device)

    # define dataloaders
    trn_loader, dev_loader, eval_loader, additional_loader, persian_loader, mp3_loader, ogg_loader, aac_loader, opus_loader = get_singfake_loaders(args.seed, config)

    # evaluates pretrained model and exit script
    if args.eval:
        model.load_state_dict(
            torch.load(config["model_path"], map_location=device))
        print("Model loaded : {}".format(config["model_path"]))
        print("Evaluating on SingFake Eval Set...")
        train_eer = evaluate(trn_loader, model, device)
        test_eer = evaluate(eval_loader, model, device)
        additional_set_eer = evaluate(additional_loader, model, device)
        persian_eer = evaluate(persian_loader, model, device)
        mp3_eer = evaluate(mp3_loader, model, device)
        ogg_eer = evaluate(ogg_loader, model, device)
        aac_eer = evaluate(aac_loader, model, device)
        opus_eer = evaluate(opus_loader, model, device)
        average_codec_eer = (mp3_eer + aac_eer + ogg_eer + opus_eer) / 4.
        print("Done. train_eer: {:.2f} %, test_set_eer: {:.2f} %, additional_test_eer: {:.2f} %, persian_eer: {:.2f} %".format(train_eer * 100, test_eer * 100, additional_set_eer * 100, persian_eer * 100))
        print("Codec testing: Average EER: {:.2f} %, MP3 EER: {:.2f} %, OGG EER: {:.2f} %, AAC EER: {:.2f} %, OPUS EER: {:.2f} %".format(average_codec_eer * 100, mp3_eer * 100, ogg_eer * 100, aac_eer * 100, opus_eer * 100))
        sys.exit(0)
        
    # get optimizer and scheduler
    optim_config["steps_per_epoch"] = len(trn_loader)
    optimizer, scheduler = create_optimizer(model.parameters(), optim_config)
    optimizer_swa = SWA(optimizer)

    best_dev_eer = 1.
    best_eval_eer = 100.
    n_swa_update = 0  # number of snapshots of model to use in SWA
    f_log = open(model_tag / "metric_log.txt", "a")
    f_log.write("=" * 5 + "\n")

    # make directory for metric logging
    metric_path = model_tag / "metrics"
    os.makedirs(metric_path, exist_ok=True)

    # Training
    for epoch in range(config["num_epochs"]):
        print("Start training epoch{:03d}".format(epoch))
        running_loss = train_epoch(trn_loader, model, optimizer, device,
                                   scheduler, config)
        
        train_eer = evaluate(trn_loader, model, device)
        dev_eer = evaluate(dev_loader, model, device)
        additional_eer = evaluate(additional_loader, model, device)
        
        print("DONE.\nLoss:{:.5f}, train_eer: {:2f} %, dev_eer: {:.2f} %, additional_test_eer: {:.2f} %".format(
            running_loss, train_eer * 100, dev_eer * 100, additional_eer * 100))
        writer.add_scalar("loss", running_loss, epoch)
        writer.add_scalar("train_eer", train_eer, epoch)
        writer.add_scalar("dev_eer", dev_eer, epoch)
        writer.add_scalar("additional_eer", additional_eer, epoch)
        
        if best_dev_eer >= dev_eer:
            print("best model find at epoch", epoch)
            best_dev_eer = dev_eer
            torch.save(model.state_dict(),
                       model_save_path / "epoch_{}_{:03.3f}.pth".format(epoch, dev_eer))

            # do evaluation whenever best model is renewed
            if str_to_bool(config["eval_all_best"]):
                eval_eer = evaluate(eval_loader, model, device)
                log_text = "epoch{:03d}, ".format(epoch)
                if eval_eer < best_eval_eer:
                    log_text += "best eer, {:.4f}%".format(eval_eer)
                    best_eval_eer = eval_eer
                    torch.save(model.state_dict(), model_save_path / "best.pth")
                    
                if len(log_text) > 0:
                    print(log_text)
                    f_log.write(log_text + "\n")

            print("Saving epoch {} for swa".format(epoch))
            optimizer_swa.update_swa()
            n_swa_update += 1
        writer.add_scalar("best_dev_eer", best_dev_eer, epoch)

    print("Start final evaluation")
    epoch += 1
    if n_swa_update > 0:
        optimizer_swa.swap_swa_sgd()
        optimizer_swa.bn_update(trn_loader, model, device=device)
    eval_eer = evaluate(eval_loader, model, device)
    f_log = open(model_tag / "metric_log.txt", "a")
    f_log.write("=" * 5 + "\n")
    f_log.write("EER: {:.3f} %".format(eval_eer * 100))
    f_log.close()

    torch.save(model.state_dict(),
               model_save_path / "swa.pth")

    if eval_eer <= best_eval_eer:
        best_eval_eer = eval_eer
        torch.save(model.state_dict(), model_save_path / "best.pth")
        
    print("Exp FIN. EER: {:.3f} %".format(
        best_eval_eer * 100))


def get_model(model_config: Dict, device: torch.device):
    """Define DNN model architecture"""
    module = import_module("models.{}".format(model_config["architecture"]))
    _model = getattr(module, "Model")
    model = _model(model_config).to(device)
    nb_params = sum([param.view(-1).size()[0] for param in model.parameters()])
    print("no. model params:{}".format(nb_params))

    return model

def get_wav2vec2_model(model_config: Dict, device: torch.device):
    model = Wav2Vec2Model(model_config, device)
    return model


def get_loader(
        database_path: str,
        seed: int,
        config: dict) -> List[torch.utils.data.DataLoader]:
    """Make PyTorch DataLoaders for train / developement / evaluation"""
    track = config["track"]
    prefix_2019 = "ASVspoof2019.{}".format(track)

    trn_database_path = database_path / "ASVspoof2019_{}_train/".format(track)
    dev_database_path = database_path / "ASVspoof2019_{}_dev/".format(track)
    eval_database_path = database_path / "ASVspoof2019_{}_eval/".format(track)

    trn_list_path = (database_path /
                     "ASVspoof2019_{}_cm_protocols/{}.cm.train.trn.txt".format(
                         track, prefix_2019))
    dev_trial_path = (database_path /
                      "ASVspoof2019_{}_cm_protocols/{}.cm.dev.trl.txt".format(
                          track, prefix_2019))
    eval_trial_path = (
        database_path /
        "ASVspoof2019_{}_cm_protocols/{}.cm.eval.trl.txt".format(
            track, prefix_2019))

    d_label_trn, file_train = genSpoof_list(dir_meta=trn_list_path,
                                            is_train=True,
                                            is_eval=False)
    print("no. training files:", len(file_train))

    train_set = Dataset_ASVspoof2019_train(list_IDs=file_train,
                                           labels=d_label_trn,
                                           base_dir=trn_database_path)
    gen = torch.Generator()
    gen.manual_seed(seed)
    trn_loader = DataLoader(train_set,
                            batch_size=config["batch_size"],
                            shuffle=True,
                            drop_last=True,
                            pin_memory=True,
                            worker_init_fn=seed_worker,
                            generator=gen)

    _, file_dev = genSpoof_list(dir_meta=dev_trial_path,
                                is_train=False,
                                is_eval=False)
    print("no. validation files:", len(file_dev))

    dev_set = Dataset_ASVspoof2019_devNeval(list_IDs=file_dev,
                                            base_dir=dev_database_path)
    dev_loader = DataLoader(dev_set,
                            batch_size=config["batch_size"],
                            shuffle=False,
                            drop_last=False,
                            pin_memory=True)

    file_eval = genSpoof_list(dir_meta=eval_trial_path,
                              is_train=False,
                              is_eval=True)
    eval_set = Dataset_ASVspoof2019_devNeval(list_IDs=file_eval,
                                             base_dir=eval_database_path)
    eval_loader = DataLoader(eval_set,
                             batch_size=config["batch_size"],
                             shuffle=False,
                             drop_last=False,
                             pin_memory=True)

    return trn_loader, dev_loader, eval_loader

def get_singfake_loaders(seed: int, config: dict) -> List[torch.utils.data.DataLoader]:
    """
    Work in progress; no split is provided as of now.
    """
    base_dir = "/home/yongyi/split_0831"
    
    target_sr = float(config["target_sr"])
    vocals_only = str_to_bool(config["vocals_only"])
    
    is_mixture = not vocals_only
    
    train_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "train"), is_mixture=is_mixture, target_sr=target_sr)
    gen = torch.Generator()
    gen.manual_seed(seed)
    trn_loader = DataLoader(train_set,
                            batch_size=config["batch_size"],
                            shuffle=True,
                            drop_last=True,
                            pin_memory=True,
                            worker_init_fn=seed_worker,
                            num_workers=4,
                            generator=gen)
    
    dev_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "dev"), is_mixture=is_mixture, target_sr=target_sr)
    dev_loader = DataLoader(dev_set,
                            batch_size=config["batch_size"],
                            shuffle=False,
                            drop_last=False,
                            num_workers=4,
                            pin_memory=True)
    
    eval_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "test"), is_mixture=is_mixture, target_sr=target_sr)
    eval_loader = DataLoader(eval_set,
                             batch_size=config["batch_size"],
                             shuffle=False,
                             drop_last=False,
                             num_workers=4,
                             pin_memory=True)
    
    additional_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "additional_test"), is_mixture=is_mixture, target_sr=target_sr)
    additional_loader = DataLoader(additional_set,
                                      batch_size=config["batch_size"],
                                      shuffle=False,
                                      drop_last=False,
                                      num_workers=4,
                                      pin_memory=True)
    
    persian_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "persian_test"), is_mixture=is_mixture, target_sr=target_sr)
    persian_loader = DataLoader(persian_set,
                                batch_size=config["batch_size"],
                                      shuffle=False,
                                      drop_last=False,
                                      num_workers=4,
                                      pin_memory=True)
    
    mp3_test_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "codec_test", "mp3_128k"), is_mixture=is_mixture, target_sr=target_sr)
    mp3_loader = DataLoader(mp3_test_set,
                            batch_size=config["batch_size"],
                                      shuffle=False,
                                      drop_last=False,
                                      num_workers=4,
                                      pin_memory=True)
    ogg_test_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "codec_test", "ogg_64k"), is_mixture=is_mixture, target_sr=target_sr)
    ogg_loader = DataLoader(ogg_test_set,
                            batch_size=config["batch_size"],
                                      shuffle=False,
                                      drop_last=False,
                                      num_workers=4,
                                      pin_memory=True)
    aac_test_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "codec_test", "adts_64k"), is_mixture=is_mixture, target_sr=target_sr)
    aac_loader = DataLoader(aac_test_set,
                            batch_size=config["batch_size"],
                                      shuffle=False,
                                      drop_last=False,
                                      num_workers=4,
                                      pin_memory=True)
    opus_test_set = Dataset_SingFake(base_dir=os.path.join(base_dir, "codec_test", "opus_64k"), is_mixture=is_mixture, target_sr=target_sr)
    opus_loader = DataLoader(opus_test_set,
                            batch_size=config["batch_size"],
                                      shuffle=False,
                                      drop_last=False,
                                      num_workers=4,
                                      pin_memory=True)
    
    return trn_loader, dev_loader, eval_loader, additional_loader, persian_loader, mp3_loader, ogg_loader, aac_loader, opus_loader
    


def evaluate(loader, model, device: torch.device):
    """
    Evaluate the model on the given loader, then return EER.
    """
    model.eval()
    # we save target (1) scores to target_scores, and non target (0) scores to nontarget_scores.
    target_scores = []
    nontarget_scores = []
    debug = False
    count = 0
    with torch.no_grad():
        for batch_x, batch_y in tqdm(loader, total=len(loader)):
            batch_x = batch_x.to(device)
            batch_out = model(batch_x)
            batch_score = (batch_out[:, 1]).data.cpu().numpy().ravel()
            batch_y = batch_y.data.cpu().numpy().ravel()
            for i in range(len(batch_y)):
                if batch_y[i] == 1:
                    target_scores.append(batch_score[i])
                else:
                    nontarget_scores.append(batch_score[i])
            count += 1
            if count == 10 and debug:
                break
    
    eer, _ = compute_eer(target_scores, nontarget_scores)
    return eer


def train_epoch(
    trn_loader: DataLoader,
    model,
    optim: Union[torch.optim.SGD, torch.optim.Adam],
    device: torch.device,
    scheduler: torch.optim.lr_scheduler,
    config: argparse.Namespace):
    """Train the model for one epoch"""
    running_loss = 0
    num_total = 0.0
    ii = 0
    model.train()

    # set objective (Loss) functions
    weight = torch.FloatTensor([0.1, 0.9]).to(device)
    criterion = nn.CrossEntropyLoss(weight=weight)
    pbar = tqdm(trn_loader, total=len(trn_loader))
    for batch_x, batch_y in pbar:
        batch_size = batch_x.size(0)
        num_total += batch_size
        ii += 1
        batch_x = batch_x.to(device)
        batch_y = batch_y.view(-1).type(torch.int64).to(device)
        batch_out = model(batch_x)
        batch_loss = criterion(batch_out, batch_y)
        running_loss += batch_loss.item() * batch_size
        pbar.set_description("loss: {:.5f}, running loss: {:.5f}".format(
            batch_loss.item(), running_loss / num_total))
        optim.zero_grad()
        batch_loss.backward()
        optim.step()

        if config["optim_config"]["scheduler"] in ["cosine", "keras_decay"]:
            scheduler.step()
        elif scheduler is None:
            pass
        else:
            raise ValueError("scheduler error, got:{}".format(scheduler))

    running_loss /= num_total
    return running_loss


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="ASVspoof detection system")
    parser.add_argument("--config",
                        dest="config",
                        type=str,
                        help="configuration file",
                        required=True)
    parser.add_argument(
        "--output_dir",
        dest="output_dir",
        type=str,
        help="output directory for results",
        default="./exp_result",
    )
    parser.add_argument("--seed",
                        type=int,
                        default=1234,
                        help="random seed (default: 1234)")
    parser.add_argument(
        "--eval",
        action="store_true",
        help="when this flag is given, evaluates given model and exit")
    parser.add_argument("--comment",
                        type=str,
                        default=None,
                        help="comment to describe the saved model")
    parser.add_argument("--eval_model_weights",
                        type=str,
                        default=None,
                        help="directory to the model weight file (can be also given in the config file)")
    parser.add_argument("--gpu",
                        type=int,
                        default=0,
                        help="gpu id to use (default: 0)")
    main(parser.parse_args())

models/wav2vec2+AASIST/main.py
yongyizang/SingFake
"""
Utilization functions
"""

import os
import random
import sys

import numpy as np
import torch


def str_to_bool(val):
    """Convert a string representation of truth to true (1) or false (0).
    Copied from the python implementation distutils.utils.strtobool

    True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values
    are 'n', 'no', 'f', 'false', 'off', and '0'.  Raises ValueError if
    'val' is anything else.
    >>> str_to_bool('YES')
    1
    >>> str_to_bool('FALSE')
    0
    """
    val = val.lower()
    if val in ('y', 'yes', 't', 'true', 'on', '1'):
        return True
    if val in ('n', 'no', 'f', 'false', 'off', '0'):
        return False
    raise ValueError('invalid truth value {}'.format(val))


def cosine_annealing(step, total_steps, lr_max, lr_min):
    """Cosine Annealing for learning rate decay scheduler"""
    return lr_min + (lr_max -
                     lr_min) * 0.5 * (1 + np.cos(step / total_steps * np.pi))


def keras_decay(step, decay=0.0001):
    """Learning rate decay in Keras-style"""
    return 1. / (1. + decay * step)


class SGDRScheduler(torch.optim.lr_scheduler._LRScheduler):
    """SGD with restarts scheduler"""
    def __init__(self, optimizer, T0, T_mul, eta_min, last_epoch=-1):
        self.Ti = T0
        self.T_mul = T_mul
        self.eta_min = eta_min

        self.last_restart = 0

        super().__init__(optimizer, last_epoch)

    def get_lr(self):
        T_cur = self.last_epoch - self.last_restart
        if T_cur >= self.Ti:
            self.last_restart = self.last_epoch
            self.Ti = self.Ti * self.T_mul
            T_cur = 0

        return [
            self.eta_min + (base_lr - self.eta_min) *
            (1 + np.cos(np.pi * T_cur / self.Ti)) / 2
            for base_lr in self.base_lrs
        ]


def _get_optimizer(model_parameters, optim_config):
    """Defines optimizer according to the given config"""
    optimizer_name = optim_config['optimizer']

    if optimizer_name == 'sgd':
        optimizer = torch.optim.SGD(model_parameters,
                                    lr=optim_config['base_lr'],
                                    momentum=optim_config['momentum'],
                                    weight_decay=optim_config['weight_decay'],
                                    nesterov=optim_config['nesterov'])
    elif optimizer_name == 'adam':
        optimizer = torch.optim.Adam(model_parameters,
                                     lr=optim_config['base_lr'],
                                     betas=optim_config['betas'],
                                     weight_decay=optim_config['weight_decay'],
                                     amsgrad=str_to_bool(
                                         optim_config['amsgrad']))
    else:
        print('Un-known optimizer', optimizer_name)
        sys.exit()

    return optimizer


def _get_scheduler(optimizer, optim_config):
    """
    Defines learning rate scheduler according to the given config
    """
    if optim_config['scheduler'] == 'multistep':
        scheduler = torch.optim.lr_scheduler.MultiStepLR(
            optimizer,
            milestones=optim_config['milestones'],
            gamma=optim_config['lr_decay'])

    elif optim_config['scheduler'] == 'sgdr':
        scheduler = SGDRScheduler(optimizer, optim_config['T0'],
                                  optim_config['Tmult'],
                                  optim_config['lr_min'])

    elif optim_config['scheduler'] == 'cosine':
        total_steps = optim_config['epochs'] * \
            optim_config['steps_per_epoch']

        scheduler = torch.optim.lr_scheduler.LambdaLR(
            optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                total_steps,
                1,  # since lr_lambda computes multiplicative factor
                optim_config['lr_min'] / optim_config['base_lr']))

    elif optim_config['scheduler'] == 'keras_decay':
        scheduler = torch.optim.lr_scheduler.LambdaLR(
            optimizer, lr_lambda=lambda step: keras_decay(step))
    else:
        scheduler = None
    return scheduler


def create_optimizer(model_parameters, optim_config):
    """Defines an optimizer and a scheduler"""
    optimizer = _get_optimizer(model_parameters, optim_config)
    scheduler = _get_scheduler(optimizer, optim_config)
    return optimizer, scheduler


def seed_worker(worker_id):
    """
    Used in generating seed for the worker of torch.utils.data.Dataloader
    """
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)


def set_seed(seed, config = None):
    """ 
    set initial seed for reproduction
    """
    if config is None:
        raise ValueError("config should not be None")

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = str_to_bool(config["cudnn_deterministic_toggle"])
        torch.backends.cudnn.benchmark = str_to_bool(config["cudnn_benchmark_toggle"])

models/wav2vec2+AASIST/utils.py
yongyizang/SingFake
import random
from typing import Union

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
import fairseq


___author__ = "Hemlata Tak"
__email__ = "tak@eurecom.fr"

############################
## FOR fine-tuned SSL MODEL
############################


class SSLModel(nn.Module):
    def __init__(self,device):
        super(SSLModel, self).__init__()
        
        cp_path = '/home/yongyi/SVspoof_dev/model/aasist/models/xlsr2_300m.pt'   # Change the pre-trained XLSR model path. 
        model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task([cp_path])
        self.model = model[0]
        self.device = device
        self.out_dim = 1024
        return

    def extract_feat(self, input_data):
        # put the model to GPU if it not there
        if next(self.model.parameters()).device != input_data.device \
           or next(self.model.parameters()).dtype != input_data.dtype:
            self.model.to(input_data.device, dtype=input_data.dtype)
            self.model.train()
        
        if True:
            # input should be in shape (batch, length)
            if input_data.ndim == 3:
                input_tmp = input_data[:, :, 0]
            else:
                input_tmp = input_data
                
            # [batch, length, dim]
            emb = self.model(input_tmp, mask=False, features_only=True)['x']
        return emb


#---------AASIST back-end------------------------#
''' Jee-weon Jung, Hee-Soo Heo, Hemlata Tak, Hye-jin Shim, Joon Son Chung, Bong-Jin Lee, Ha-Jin Yu and Nicholas Evans. 
    AASIST: Audio Anti-Spoofing Using Integrated Spectro-Temporal Graph Attention Networks. 
    In Proc. ICASSP 2022, pp: 6367--6371.'''


class GraphAttentionLayer(nn.Module):
    def __init__(self, in_dim, out_dim, **kwargs):
        super().__init__()

        # attention map
        self.att_proj = nn.Linear(in_dim, out_dim)
        self.att_weight = self._init_new_params(out_dim, 1)

        # project
        self.proj_with_att = nn.Linear(in_dim, out_dim)
        self.proj_without_att = nn.Linear(in_dim, out_dim)

        # batch norm
        self.bn = nn.BatchNorm1d(out_dim)

        # dropout for inputs
        self.input_drop = nn.Dropout(p=0.2)

        # activate
        self.act = nn.SELU(inplace=True)

        # temperature
        self.temp = 1.
        if "temperature" in kwargs:
            self.temp = kwargs["temperature"]

    def forward(self, x):
        '''
        x   :(#bs, #node, #dim)
        '''
        # apply input dropout
        x = self.input_drop(x)

        # derive attention map
        att_map = self._derive_att_map(x)

        # projection
        x = self._project(x, att_map)

        # apply batch norm
        x = self._apply_BN(x)
        x = self.act(x)
        return x

    def _pairwise_mul_nodes(self, x):
        '''
        Calculates pairwise multiplication of nodes.
        - for attention map
        x           :(#bs, #node, #dim)
        out_shape   :(#bs, #node, #node, #dim)
        '''

        nb_nodes = x.size(1)
        x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1)
        x_mirror = x.transpose(1, 2)

        return x * x_mirror

    def _derive_att_map(self, x):
        '''
        x           :(#bs, #node, #dim)
        out_shape   :(#bs, #node, #node, 1)
        '''
        att_map = self._pairwise_mul_nodes(x)
        # size: (#bs, #node, #node, #dim_out)
        att_map = torch.tanh(self.att_proj(att_map))
        # size: (#bs, #node, #node, 1)
        att_map = torch.matmul(att_map, self.att_weight)

        # apply temperature
        att_map = att_map / self.temp

        att_map = F.softmax(att_map, dim=-2)

        return att_map

    def _project(self, x, att_map):
        x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x))
        x2 = self.proj_without_att(x)

        return x1 + x2

    def _apply_BN(self, x):
        org_size = x.size()
        x = x.view(-1, org_size[-1])
        x = self.bn(x)
        x = x.view(org_size)

        return x

    def _init_new_params(self, *size):
        out = nn.Parameter(torch.FloatTensor(*size))
        nn.init.xavier_normal_(out)
        return out


class HtrgGraphAttentionLayer(nn.Module):
    def __init__(self, in_dim, out_dim, **kwargs):
        super().__init__()

        self.proj_type1 = nn.Linear(in_dim, in_dim)
        self.proj_type2 = nn.Linear(in_dim, in_dim)

        # attention map
        self.att_proj = nn.Linear(in_dim, out_dim)
        self.att_projM = nn.Linear(in_dim, out_dim)

        self.att_weight11 = self._init_new_params(out_dim, 1)
        self.att_weight22 = self._init_new_params(out_dim, 1)
        self.att_weight12 = self._init_new_params(out_dim, 1)
        self.att_weightM = self._init_new_params(out_dim, 1)

        # project
        self.proj_with_att = nn.Linear(in_dim, out_dim)
        self.proj_without_att = nn.Linear(in_dim, out_dim)

        self.proj_with_attM = nn.Linear(in_dim, out_dim)
        self.proj_without_attM = nn.Linear(in_dim, out_dim)

        # batch norm
        self.bn = nn.BatchNorm1d(out_dim)

        # dropout for inputs
        self.input_drop = nn.Dropout(p=0.2)

        # activate
        self.act = nn.SELU(inplace=True)

        # temperature
        self.temp = 1.
        if "temperature" in kwargs:
            self.temp = kwargs["temperature"]

    def forward(self, x1, x2, master=None):
        '''
        x1  :(#bs, #node, #dim)
        x2  :(#bs, #node, #dim)
        '''
        #print('x1',x1.shape)
        #print('x2',x2.shape)
        num_type1 = x1.size(1)
        num_type2 = x2.size(1)
        #print('num_type1',num_type1)
        #print('num_type2',num_type2)
        x1 = self.proj_type1(x1)
        #print('proj_type1',x1.shape)
        x2 = self.proj_type2(x2)
        #print('proj_type2',x2.shape)
        x = torch.cat([x1, x2], dim=1)
        #print('Concat x1 and x2',x.shape)
        
        if master is None:
            master = torch.mean(x, dim=1, keepdim=True)
            #print('master',master.shape)
        # apply input dropout
        x = self.input_drop(x)

        # derive attention map
        att_map = self._derive_att_map(x, num_type1, num_type2)
        #print('master',master.shape)
        # directional edge for master node
        master = self._update_master(x, master)
        #print('master',master.shape)
        # projection
        x = self._project(x, att_map)
        #print('proj x',x.shape)
        # apply batch norm
        x = self._apply_BN(x)
        x = self.act(x)

        x1 = x.narrow(1, 0, num_type1)
        #print('x1',x1.shape)
        x2 = x.narrow(1, num_type1, num_type2)
        #print('x2',x2.shape)
        return x1, x2, master

    def _update_master(self, x, master):

        att_map = self._derive_att_map_master(x, master)
        master = self._project_master(x, master, att_map)

        return master

    def _pairwise_mul_nodes(self, x):
        '''
        Calculates pairwise multiplication of nodes.
        - for attention map
        x           :(#bs, #node, #dim)
        out_shape   :(#bs, #node, #node, #dim)
        '''

        nb_nodes = x.size(1)
        x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1)
        x_mirror = x.transpose(1, 2)

        return x * x_mirror

    def _derive_att_map_master(self, x, master):
        '''
        x           :(#bs, #node, #dim)
        out_shape   :(#bs, #node, #node, 1)
        '''
        att_map = x * master
        att_map = torch.tanh(self.att_projM(att_map))

        att_map = torch.matmul(att_map, self.att_weightM)

        # apply temperature
        att_map = att_map / self.temp

        att_map = F.softmax(att_map, dim=-2)

        return att_map

    def _derive_att_map(self, x, num_type1, num_type2):
        '''
        x           :(#bs, #node, #dim)
        out_shape   :(#bs, #node, #node, 1)
        '''
        att_map = self._pairwise_mul_nodes(x)
        # size: (#bs, #node, #node, #dim_out)
        att_map = torch.tanh(self.att_proj(att_map))
        # size: (#bs, #node, #node, 1)

        att_board = torch.zeros_like(att_map[:, :, :, 0]).unsqueeze(-1)

        att_board[:, :num_type1, :num_type1, :] = torch.matmul(
            att_map[:, :num_type1, :num_type1, :], self.att_weight11)
        att_board[:, num_type1:, num_type1:, :] = torch.matmul(
            att_map[:, num_type1:, num_type1:, :], self.att_weight22)
        att_board[:, :num_type1, num_type1:, :] = torch.matmul(
            att_map[:, :num_type1, num_type1:, :], self.att_weight12)
        att_board[:, num_type1:, :num_type1, :] = torch.matmul(
            att_map[:, num_type1:, :num_type1, :], self.att_weight12)

        att_map = att_board

        

        # apply temperature
        att_map = att_map / self.temp

        att_map = F.softmax(att_map, dim=-2)

        return att_map

    def _project(self, x, att_map):
        x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x))
        x2 = self.proj_without_att(x)

        return x1 + x2

    def _project_master(self, x, master, att_map):

        x1 = self.proj_with_attM(torch.matmul(
            att_map.squeeze(-1).unsqueeze(1), x))
        x2 = self.proj_without_attM(master)

        return x1 + x2

    def _apply_BN(self, x):
        org_size = x.size()
        x = x.view(-1, org_size[-1])
        x = self.bn(x)
        x = x.view(org_size)

        return x

    def _init_new_params(self, *size):
        out = nn.Parameter(torch.FloatTensor(*size))
        nn.init.xavier_normal_(out)
        return out


class GraphPool(nn.Module):
    def __init__(self, k: float, in_dim: int, p: Union[float, int]):
        super().__init__()
        self.k = k
        self.sigmoid = nn.Sigmoid()
        self.proj = nn.Linear(in_dim, 1)
        self.drop = nn.Dropout(p=p) if p > 0 else nn.Identity()
        self.in_dim = in_dim

    def forward(self, h):
        Z = self.drop(h)
        weights = self.proj(Z)
        scores = self.sigmoid(weights)
        new_h = self.top_k_graph(scores, h, self.k)

        return new_h

    def top_k_graph(self, scores, h, k):
        """
        args
        =====
        scores: attention-based weights (#bs, #node, 1)
        h: graph data (#bs, #node, #dim)
        k: ratio of remaining nodes, (float)
        returns
        =====
        h: graph pool applied data (#bs, #node', #dim)
        """
        _, n_nodes, n_feat = h.size()
        n_nodes = max(int(n_nodes * k), 1)
        _, idx = torch.topk(scores, n_nodes, dim=1)
        idx = idx.expand(-1, -1, n_feat)

        h = h * scores
        h = torch.gather(h, 1, idx)

        return h




class Residual_block(nn.Module):
    def __init__(self, nb_filts, first=False):
        super().__init__()
        self.first = first

        if not self.first:
            self.bn1 = nn.BatchNorm2d(num_features=nb_filts[0])
        self.conv1 = nn.Conv2d(in_channels=nb_filts[0],
                               out_channels=nb_filts[1],
                               kernel_size=(2, 3),
                               padding=(1, 1),
                               stride=1)
        self.selu = nn.SELU(inplace=True)

        self.bn2 = nn.BatchNorm2d(num_features=nb_filts[1])
        self.conv2 = nn.Conv2d(in_channels=nb_filts[1],
                               out_channels=nb_filts[1],
                               kernel_size=(2, 3),
                               padding=(0, 1),
                               stride=1)

        if nb_filts[0] != nb_filts[1]:
            self.downsample = True
            self.conv_downsample = nn.Conv2d(in_channels=nb_filts[0],
                                             out_channels=nb_filts[1],
                                             padding=(0, 1),
                                             kernel_size=(1, 3),
                                             stride=1)

        else:
            self.downsample = False
        

    def forward(self, x):
        identity = x
        if not self.first:
            out = self.bn1(x)
            out = self.selu(out)
        else:
            out = x

        #print('out',out.shape)
        out = self.conv1(x)

        #print('aft conv1 out',out.shape)
        out = self.bn2(out)
        out = self.selu(out)
        # print('out',out.shape)
        out = self.conv2(out)
        #print('conv2 out',out.shape)
        
        if self.downsample:
            identity = self.conv_downsample(identity)

        out += identity
        #out = self.mp(out)
        return out


class Wav2Vec2Model(nn.Module):
    def __init__(self, args,device):
        super().__init__()
        self.device = device
        
        # AASIST parameters
        filts = [128, [1, 32], [32, 32], [32, 64], [64, 64]]
        gat_dims = [64, 32]
        pool_ratios = [0.5, 0.5, 0.5, 0.5]
        temperatures =  [2.0, 2.0, 100.0, 100.0]


        ####
        # create network wav2vec 2.0
        ####
        self.ssl_model = SSLModel(self.device)
        self.LL = nn.Linear(self.ssl_model.out_dim, 128)

        self.first_bn = nn.BatchNorm2d(num_features=1)
        self.first_bn1 = nn.BatchNorm2d(num_features=64)
        self.drop = nn.Dropout(0.5, inplace=True)
        self.drop_way = nn.Dropout(0.2, inplace=True)
        self.selu = nn.SELU(inplace=True)

        # RawNet2 encoder
        self.encoder = nn.Sequential(
            nn.Sequential(Residual_block(nb_filts=filts[1], first=True)),
            nn.Sequential(Residual_block(nb_filts=filts[2])),
            nn.Sequential(Residual_block(nb_filts=filts[3])),
            nn.Sequential(Residual_block(nb_filts=filts[4])),
            nn.Sequential(Residual_block(nb_filts=filts[4])),
            nn.Sequential(Residual_block(nb_filts=filts[4])))

        self.attention = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=(1,1)),
            nn.SELU(inplace=True),
            nn.BatchNorm2d(128),
            nn.Conv2d(128, 64, kernel_size=(1,1)),
            
        )
        # position encoding
        self.pos_S = nn.Parameter(torch.randn(1, 42, filts[-1][-1]))
        
        self.master1 = nn.Parameter(torch.randn(1, 1, gat_dims[0]))
        self.master2 = nn.Parameter(torch.randn(1, 1, gat_dims[0]))
        
        # Graph module
        self.GAT_layer_S = GraphAttentionLayer(filts[-1][-1],
                                               gat_dims[0],
                                               temperature=temperatures[0])
        self.GAT_layer_T = GraphAttentionLayer(filts[-1][-1],
                                               gat_dims[0],
                                               temperature=temperatures[1])
        # HS-GAL layer 
        self.HtrgGAT_layer_ST11 = HtrgGraphAttentionLayer(
            gat_dims[0], gat_dims[1], temperature=temperatures[2])
        self.HtrgGAT_layer_ST12 = HtrgGraphAttentionLayer(
            gat_dims[1], gat_dims[1], temperature=temperatures[2])
        self.HtrgGAT_layer_ST21 = HtrgGraphAttentionLayer(
            gat_dims[0], gat_dims[1], temperature=temperatures[2])
        self.HtrgGAT_layer_ST22 = HtrgGraphAttentionLayer(
            gat_dims[1], gat_dims[1], temperature=temperatures[2])

        # Graph pooling layers
        self.pool_S = GraphPool(pool_ratios[0], gat_dims[0], 0.3)
        self.pool_T = GraphPool(pool_ratios[1], gat_dims[0], 0.3)
        self.pool_hS1 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)
        self.pool_hT1 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)

        self.pool_hS2 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)
        self.pool_hT2 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)
        
        self.out_layer = nn.Linear(5 * gat_dims[1], 2)

    def forward(self, x):
        #-------pre-trained Wav2vec model fine tunning ------------------------##
        x_ssl_feat = self.ssl_model.extract_feat(x.squeeze(-1))
        x = self.LL(x_ssl_feat) #(bs,frame_number,feat_out_dim)
        
        # post-processing on front-end features
        x = x.transpose(1, 2)   #(bs,feat_out_dim,frame_number)
        x = x.unsqueeze(dim=1) # add channel 
        x = F.max_pool2d(x, (3, 3))
        x = self.first_bn(x)
        x = self.selu(x)

        # RawNet2-based encoder
        x = self.encoder(x)
        x = self.first_bn1(x)
        x = self.selu(x)
        
        w = self.attention(x)
        
        #------------SA for spectral feature-------------#
        w1 = F.softmax(w,dim=-1)
        m = torch.sum(x * w1, dim=-1)
        e_S = m.transpose(1, 2) + self.pos_S 
        
        # graph module layer
        gat_S = self.GAT_layer_S(e_S)
        out_S = self.pool_S(gat_S)  # (#bs, #node, #dim)
        
        #------------SA for temporal feature-------------#
        w2 = F.softmax(w,dim=-2)
        m1 = torch.sum(x * w2, dim=-2)
     
        e_T = m1.transpose(1, 2)
       
        # graph module layer
        gat_T = self.GAT_layer_T(e_T)
        out_T = self.pool_T(gat_T)
        
        # learnable master node
        master1 = self.master1.expand(x.size(0), -1, -1)
        master2 = self.master2.expand(x.size(0), -1, -1)

        # inference 1
        out_T1, out_S1, master1 = self.HtrgGAT_layer_ST11(
            out_T, out_S, master=self.master1)

        out_S1 = self.pool_hS1(out_S1)
        out_T1 = self.pool_hT1(out_T1)

        out_T_aug, out_S_aug, master_aug = self.HtrgGAT_layer_ST12(
            out_T1, out_S1, master=master1)
        out_T1 = out_T1 + out_T_aug
        out_S1 = out_S1 + out_S_aug
        master1 = master1 + master_aug

        # inference 2
        out_T2, out_S2, master2 = self.HtrgGAT_layer_ST21(
            out_T, out_S, master=self.master2)
        out_S2 = self.pool_hS2(out_S2)
        out_T2 = self.pool_hT2(out_T2)

        out_T_aug, out_S_aug, master_aug = self.HtrgGAT_layer_ST22(
            out_T2, out_S2, master=master2)
        out_T2 = out_T2 + out_T_aug
        out_S2 = out_S2 + out_S_aug
        master2 = master2 + master_aug

        out_T1 = self.drop_way(out_T1)
        out_T2 = self.drop_way(out_T2)
        out_S1 = self.drop_way(out_S1)
        out_S2 = self.drop_way(out_S2)
        master1 = self.drop_way(master1)
        master2 = self.drop_way(master2)

        out_T = torch.max(out_T1, out_T2)
        out_S = torch.max(out_S1, out_S2)
        master = torch.max(master1, master2)

        # Readout operation
        T_max, _ = torch.max(torch.abs(out_T), dim=1)
        T_avg = torch.mean(out_T, dim=1)

        S_max, _ = torch.max(torch.abs(out_S), dim=1)
        S_avg = torch.mean(out_S, dim=1)
        
        last_hidden = torch.cat(
            [T_max, T_avg, S_max, S_avg, master.squeeze(1)], dim=1)
        
        last_hidden = self.drop(last_hidden)
        output = self.out_layer(last_hidden)
        
        return output

models/wav2vec2+AASIST/wav2vecAASIST.py
yongyizang/SingFake